fleet/cmd/fleetctl/scripts.go

package main

import (
	"errors"
	"fmt"
	"html/template"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"unicode/utf8"

	"github.com/fleetdm/fleet/v4/server/fleet"
	"github.com/fleetdm/fleet/v4/server/service"
	"github.com/urfave/cli/v2"
)

func runScriptCommand() *cli.Command {
	return &cli.Command{
		Name:      "run-script",
		Aliases:   []string{"run_script"},
		Usage:     `Run a live script on one host and get results back (5 minute timeout).`,
		UsageText: `fleetctl run-script [options]`,
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "script-path",
				Usage:    "The path to the script.",
				Required: false,
			},
			&cli.StringFlag{
				Name:     "host",
				Usage:    "A host, specified by hostname, serial number, UUID, osquery host ID, or node key.",
				Required: true,
			},
			&cli.StringFlag{
				Name:     "script-name",
				Usage:    "Name of saved script to run.",
				Required: false,
			},
			&cli.UintFlag{
				Name:     "team-id",
				Usage:    `Available in Fleet Premium. ID of the team that the saved script belongs to. 0 targets hosts assigned to “No team” (default: 0).`,
				Required: false,
			},
			configFlag(),
			contextFlag(),
			debugFlag(),
		},
		Action: func(c *cli.Context) error {
			client, err := clientFromCLI(c)
			if err != nil {
				return err
			}

			appCfg, err := client.GetAppConfig()
			if err != nil {
				return err
			}

			if appCfg.ServerSettings.ScriptsDisabled {
				return errors.New(fleet.RunScriptScriptsDisabledGloballyErrMsg)
			}

			path := c.String("script-path")
			name := c.String("script-name")

			if path == "" && name == "" {
				return errors.New("One of --script-path or --script-name must be specified.")
			}

			if path != "" && name != "" {
				return errors.New("Only one of --script-path or --script-name is allowed.")
			}

			if path != "" {
				if err := validateScriptPath(path); err != nil {
					return err
				}
			}

			ident := c.String("host")
			h, err := client.HostByIdentifier(ident)
			if err != nil {
				var nfe service.NotFoundErr
				if errors.As(err, &nfe) {
					return errors.New(fleet.RunScriptHostNotFoundErrMsg)
				}
				var sce fleet.ErrWithStatusCode
				if errors.As(err, &sce) {
					if sce.StatusCode() == http.StatusForbidden {
						return errors.New(fleet.RunScriptForbiddenErrMsg)
					}
				}
				return err
			}

			if h.Status != fleet.StatusOnline {
				return errors.New(fleet.RunScriptHostOfflineErrMsg)
			}

			var b []byte
			if path != "" {
				b, err = os.ReadFile(path)
				if err != nil {
					return err
				}

				// validate script contents with isSavedScript flag set to false so that we check
				// for the shorter
				if err := fleet.ValidateHostScriptContents(string(b), false); err != nil {
					if err.Error() == fleet.RunScripUnsavedMaxLenErrMsg {
						return errors.New("Script is too large. Script referenced by '--script-path' is limited to 10,000 characters. To run larger script save it to Fleet and use '--script-name'.")
					}
					return err
				}
			}

			fmt.Println("\nScript is running. Please wait for it to finish...")

			res, err := client.RunHostScriptSync(h.ID, b, name, c.Uint("team-id"))
			if err != nil {
				if strings.Contains(err.Error(), `Only one of 'script_contents' or 'team_id' is allowed`) {
					return errors.New("Only one of --script-path or --team-id is allowed.")
				}
				return err
			}

			if err := renderScriptResult(c, res); err != nil {
				return err
			}

			return nil
		},
	}
}

func renderScriptResult(c *cli.Context, res *fleet.HostScriptResult) error {
	tmpl := template.Must(template.New("").Parse(`
{{ if .ErrorMsg -}}
Error: {{ .ErrorMsg }}
{{- else -}}
Exit code: {{ .ExitCode }} ({{ .ExitMessage }})
{{- end }}
{{ if .ShowOutput }}
Output {{- if .ExecTimeout }} before timeout {{- end }}:

-------------------------------------------------------------------------------------

{{ .Output }}

-------------------------------------------------------------------------------------
{{- end }}
`))

	data := struct {
		ExecTimeout bool
		ErrorMsg    string
		ExitCode    *int64
		ExitMessage string
		Output      string
		ShowOutput  bool
	}{
		ExitCode:    res.ExitCode,
		ExitMessage: "Script failed.",
		ShowOutput:  true,
	}

	switch {
	case res.ExitCode == nil:
		data.ErrorMsg = res.Message
	case *res.ExitCode == -2:
		data.ShowOutput = false
		data.ErrorMsg = res.Message
	case *res.ExitCode == -1:
		data.ExecTimeout = true
		data.ErrorMsg = res.Message
	case *res.ExitCode == 0:
		data.ExitMessage = "Script ran successfully."
	}

	if len(res.Output) >= fleet.UnsavedScriptMaxRuneLen && utf8.RuneCountInString(res.Output) >= fleet.UnsavedScriptMaxRuneLen {
		data.Output = "Fleet records the last 10,000 characters to prevent downtime.\n\n" + res.Output
	} else {
		data.Output = res.Output
	}

	return tmpl.Execute(c.App.Writer, data)
}

func validateScriptPath(path string) error {
	extension := filepath.Ext(path)
	if extension == ".sh" || extension == ".ps1" {
		return nil
	}
	return errors.New(fleet.RunScriptInvalidTypeErrMsg)
}