fleet/orbit/pkg/setup_experience/setup_experience.go

package setupexperience

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/fleetdm/fleet/v4/orbit/pkg/constant"
	"github.com/fleetdm/fleet/v4/orbit/pkg/swiftdialog"
	"github.com/fleetdm/fleet/v4/orbit/pkg/token"
	"github.com/fleetdm/fleet/v4/orbit/pkg/update"
	"github.com/fleetdm/fleet/v4/server/fleet"
	"github.com/fleetdm/fleet/v4/server/ptr"
	"github.com/rs/zerolog/log"
)

// OrbitClient is the minimal interface needed to communicate with the Fleet server.
type OrbitClient interface {
	GetSetupExperienceStatus(resetFailedSetupSteps bool) (*fleet.SetupExperienceStatusPayload, error)
}

// DeviceClient is the minimal interface needed to get the device's browser URL.
type DeviceClient interface {
	BrowserDeviceURL(token string) string
}

// SetupExperiencer is the type that manages the Fleet setup experience flow during macOS Setup
// Assistant. It uses swiftDialog as a UI for showing the status of software installations and
// script execution that are configured to run before the user has full access to the device.
// If the setup experience is supposed to run, it will launch a single swiftDialog instance and then
// update that instance based on the results from the /orbit/setup_experience/status endpoint.
type SetupExperiencer struct {
	OrbitClient  OrbitClient
	DeviceClient DeviceClient
	closeChan    chan struct{}
	rootDirPath  string
	// Note: this object is not safe for concurrent use. Since the SetupExperiencer is a singleton,
	// its Run method is called within a WaitGroup,
	// and no other parts of Orbit need access to this field (or any other parts of the
	// SetupExperiencer), it's OK to not protect this with a lock.
	sd                *swiftdialog.SwiftDialog
	started           bool
	trw               *token.ReadWriter
	stopTokenRotation func()
	uiSteps           map[string]swiftdialog.ListItem // For legacy swiftDialog UI to track current steps
	UseLegacyUI       bool
}

func NewSetupExperiencer(orbitClient OrbitClient, deviceClient DeviceClient, rootDirPath string, trw *token.ReadWriter) *SetupExperiencer {
	return &SetupExperiencer{
		OrbitClient:  orbitClient,
		DeviceClient: deviceClient,
		closeChan:    make(chan struct{}),
		rootDirPath:  rootDirPath,
		trw:          trw,
		uiSteps:      make(map[string]swiftdialog.ListItem),
	}
}

func (s *SetupExperiencer) Run(oc *fleet.OrbitConfig) error {
	if !oc.Notifications.RunSetupExperience {
		log.Debug().Msg("skipping setup experience: notification flag is not set")
		return nil
	}

	// If using the legacy UI, then call that method.
	if s.UseLegacyUI {
		return s.RunLegacy(oc)
	}

	// Ensure that the token rotation checker is started, so that we have a valid token
	// when we need to show or refresh the My Device URL in the webview.
	if s.stopTokenRotation == nil {
		s.stopTokenRotation = s.trw.StartRotation()
	}

	_, binaryPath, _ := update.LocalTargetPaths(
		s.rootDirPath,
		"swiftDialog",
		update.SwiftDialogMacOSTarget,
	)

	if _, err := os.Stat(binaryPath); err != nil {
		log.Info().Msg("skipping setup experience: swiftDialog is not installed")
		return nil
	}

	log.Info().Msg("checking setup experience status")

	// Poll the status endpoint. This also releases the device if we're done.
	payload, err := s.OrbitClient.GetSetupExperienceStatus(!s.started)
	if err != nil {
		return err
	}
	// Marshall the payload for logging
	payloadBytes, err := json.Marshal(payload)
	if err != nil {
		log.Error().Err(err).Msg("marshalling setup experience payload for logging")
	} else {
		log.Debug().Msgf("setup experience payload: %s", string(payloadBytes))
	}

	// If swiftDialog isn't up yet, then launch it
	orgLogo := payload.OrgLogoURL
	if orgLogo == "" {
		orgLogo = "https://fleetdm.com/images/permanent/fleet-mark-color-40x40@4x.png"
	}

	if err := s.startSwiftDialog(binaryPath, orgLogo); err != nil {
		return err
	}

	// Defer this so that s.started is only false the first time this function runs.
	defer func() { s.started = true }()

	select {
	case <-s.closeChan:
		log.Info().Str("receiver", "setup_experiencer").Msg("swiftDialog closed")
		return nil
	default:
		// ok
	}

	// We're rendering the initial loading UI (shown while there are still profiles, bootstrap package,
	// and account configuration to verify) right off the bat, so we can just no-op if any of those
	// are not terminal

	log.Info().Msg("setup experience: checking for pending statuses")

	if payload.BootstrapPackage != nil {
		if payload.BootstrapPackage.Status != fleet.MDMBootstrapPackageFailed && payload.BootstrapPackage.Status != fleet.MDMBootstrapPackageInstalled {
			log.Info().Msg("setup experience: bootstrap package pending")
			return nil
		}
	}

	if isPending, name := anyProfilePending(payload.ConfigurationProfiles); isPending {
		log.Info().Msg(fmt.Sprintf("setup experience: profile pending: %s", name))
		return nil
	}

	if payload.AccountConfiguration != nil {
		if payload.AccountConfiguration.Status != fleet.MDMAppleStatusAcknowledged &&
			payload.AccountConfiguration.Status != fleet.MDMAppleStatusError &&
			payload.AccountConfiguration.Status != fleet.MDMAppleStatusCommandFormatError {

			log.Info().Msg("setup experience: account config pending")
			return nil
		}
	}

	// If we got this far, then we can hand the UI over to the webview.

	// Clear the dialog message.
	if err := s.sd.HideMessage(); err != nil {
		log.Error().Err(err).Msg("clearing message in setup experience UI")
	}
	// Remove the icon.
	if err := s.sd.HideIcon(); err != nil {
		log.Error().Err(err).Msg("clearing icon in setup experience UI")
	}
	// Hide the title.
	if err := s.sd.HideTitle(); err != nil {
		log.Error().Err(err).Msg("hiding title in setup experience UI")
	}
	// Hide the progress.
	if err := s.sd.HideProgress(); err != nil {
		log.Error().Err(err).Msg("hiding progress in setup experience UI")
	}
	// Get the device token.
	token, err := s.trw.Read()
	if err != nil {
		return fmt.Errorf("getting device token: %w", err)
	}
	// Get the My Device URL.
	browserURL := s.DeviceClient.BrowserDeviceURL(token)
	// log out the url
	log.Debug().Msgf("setup experience: opening web content URL: %s", browserURL)
	// Set the web content URL.
	if err := s.sd.SetWebContent(browserURL + "?setup_only=1"); err != nil {
		log.Error().Err(err).Msg("setting web content URL in setup experience UI")
		return nil
	}

	// Note that we are setting this based on the current payload only just in case something
	// was removed from the payload that was there earlier(e.g. a deleted software title).
	allStepsDone := true

	// Now render the UI for the software and script.
	if len(payload.Software) > 0 || payload.Script != nil {
		log.Info().Msg("setup experience: rendering software and script UI")

		if len(payload.Software) > 0 {
			for _, step := range payload.Software {
				// If any step is not in a terminal state, then we're not done.
				if (step.Status != fleet.SetupExperienceStatusFailure && step.Status != fleet.SetupExperienceStatusSuccess) ||
					// If any software failed, and we're requiring all software to succeed, then we'll block completion.
					(step.Status == fleet.SetupExperienceStatusFailure && payload.RequireAllSoftware) {
					allStepsDone = false
				}
			}
		}

		// If a script is still running, then we're not done.
		if payload.Script != nil && payload.Script.Status != fleet.SetupExperienceStatusFailure && payload.Script.Status != fleet.SetupExperienceStatusSuccess {
			allStepsDone = false
		}
	}

	// If we get here, we can close the webview.
	// It will likely already be displaying a "done" message.
	if allStepsDone {
		if err := s.sd.EnableButton1(true); err != nil {
			log.Info().Err(err).Msg("enabling close button in setup experience UI")
		}

		// Sleep for a few seconds to let the user see the done message before closing
		// the UI
		time.Sleep(3 * time.Second)

		if err := s.sd.Quit(); err != nil {
			log.Info().Err(err).Msg("quitting setup experience UI on completion")
		}

		// Stop the token rotation checker since we're done with the setup experience.
		s.stopTokenRotation()
	}

	return nil
}

func anyProfilePending(profiles []*fleet.SetupExperienceConfigurationProfileResult) (bool, string) {
	for _, p := range profiles {
		if p.Status == fleet.MDMDeliveryPending {
			return true, p.Name
		}
	}

	return false, ""
}

func (s *SetupExperiencer) startSwiftDialog(binaryPath, orgLogo string) error {
	if s.started {
		log.Info().Msg("swiftDialog started")
		return nil
	}

	log.Info().Msg("creating swiftDialog instance")

	created := make(chan struct{})
	swiftDialog, err := swiftdialog.Create(context.Background(), binaryPath)
	if err != nil {
		return errors.New("creating swiftDialog instance: %w")
	}
	s.sd = swiftDialog

	iconSize, err := swiftdialog.GetIconSize(orgLogo)
	if err != nil {
		log.Error().Err(err).Msg("setup experience: getting icon size")
		iconSize = swiftdialog.DefaultIconSize
	}

	go func() {
		initOpts := &swiftdialog.SwiftDialogOptions{
			Title:            "none",
			Message:          "### Setting up your Mac...\n\nYour Mac is being configured by your organization using Fleet. This process may take some time to complete. Please don't attempt to restart or shut down the computer unless prompted to do so.",
			Icon:             orgLogo,
			MessageAlignment: swiftdialog.AlignmentCenter,
			CentreIcon:       true,
			Height:           "625",
			Big:              true,
			ProgressText:     "Configuring your device...",
			Button1Text:      "Close",
			Button1Disabled:  true,
			BlurScreen:       true,
			OnTop:            true,
			QuitKey:          "X", // Capital X to require command+shift+x
		}

		if err := s.sd.Start(context.Background(), initOpts, true); err != nil {
			log.Error().Err(err).Msg("starting swiftDialog instance")
		}

		if err = s.sd.ShowProgress(); err != nil {
			log.Error().Err(err).Msg("setting initial setup experience progress")
		}

		if err := s.sd.SetIconSize(iconSize); err != nil {
			log.Error().Err(err).Msg("setting initial setup experience icon size")
		}

		log.Debug().Msg("swiftDialog process started")
		created <- struct{}{}

		if _, err = s.sd.Wait(); err != nil {
			log.Error().Err(err).Msg("swiftdialog.Wait failed")
		}

		s.closeChan <- struct{}{}
	}()
	<-created
	return nil
}

// LinuxSetupExperiencer runs the setup experience on Linux hosts.
type LinuxSetupExperiencer struct {
	orbitClient OrbitClient
	rootDir     string
}

// NewLinuxSetupExperiencer creates a config receiver to run the setup experience on Linux hosts.
func NewLinuxSetupExperiencer(client OrbitClient, rootDir string) *LinuxSetupExperiencer {
	return &LinuxSetupExperiencer{
		orbitClient: client,
		rootDir:     rootDir,
	}
}

// Run implements fleet.OrbitConfigReceiver.
//
// Currently the fleet.OrbitConfig is ununsed but might be used in the future.
func (s *LinuxSetupExperiencer) Run(_ *fleet.OrbitConfig) error {
	info, err := ReadSetupExperienceStatusFile(s.rootDir)
	if err != nil {
		return fmt.Errorf("read setup experience file: %w", err)
	}
	if info == nil || info.TimeFinished != nil {
		// nothing to do.
		return nil
	}

	payload, err := s.orbitClient.GetSetupExperienceStatus(false)
	if err != nil {
		return err
	}

	if setupExperienceDone(payload) {
		info.TimeFinished = ptr.Time(time.Now())
		if err := WriteSetupExperienceStatusFile(s.rootDir, info); err != nil {
			log.Error().Err(err).Msg("write setup experience status file")
		}
	}
	return nil
}

func setupExperienceDone(payload *fleet.SetupExperienceStatusPayload) bool {
	for _, software := range payload.Software {
		if software != nil && (software.Status == fleet.SetupExperienceStatusPending || software.Status == fleet.SetupExperienceStatusRunning) {
			return false
		}
	}
	return true
}

// SetupExperienceInfo holds information of the state of the setup experience for a host.
type SetupExperienceInfo struct {
	// TimeInitiated is the time the setup experience was attempted during setup/installation.
	TimeInitiated time.Time `json:"time_initiated"`
	// Enabled is true if the setup experience was enabled during setup/installation.
	Enabled bool `json:"enabled"`
	// TimeFinished is the time the setup experience was finished by the host.
	TimeFinished *time.Time `json:"time_finished,omitempty"`
}

// ReadSetupExperienceStatusFile reads the setup experience state from a known file in the rootDir.
func ReadSetupExperienceStatusFile(rootDir string) (*SetupExperienceInfo, error) {
	infoPath := filepath.Join(rootDir, constant.SetupExperienceFilename)

	f, err := os.Open(infoPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("read setup experience file: %w", err)
	}
	defer f.Close()
	var exp SetupExperienceInfo
	if err := json.NewDecoder(f).Decode(&exp); err != nil {
		return nil, fmt.Errorf("decoding setup experience file: %w", err)
	}

	return &exp, nil
}

// WriteSetupExperienceStatusFile writes the setup experience state to a file under rootDir.
func WriteSetupExperienceStatusFile(rootDir string, exp *SetupExperienceInfo) error {
	infoPath := filepath.Join(rootDir, constant.SetupExperienceFilename)

	f, err := os.OpenFile(infoPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, constant.DefaultFileMode)
	if err != nil {
		return fmt.Errorf("create setup experience completed file: %w", err)
	}
	defer f.Close()
	if err := json.NewEncoder(f).Encode(exp); err != nil {
		return fmt.Errorf("write setup experience completed file: %w", err)
	}
	return nil
}

func (s *SetupExperiencer) RunLegacy(oc *fleet.OrbitConfig) error {
	const doneMessage = `### Setup is complete\n\nPlease contact your IT Administrator if there were any errors.`

	if !oc.Notifications.RunSetupExperience {
		log.Debug().Msg("skipping setup experience: notification flag is not set")
		return nil
	}

	_, binaryPath, _ := update.LocalTargetPaths(
		s.rootDirPath,
		"swiftDialog",
		update.SwiftDialogMacOSTarget,
	)

	if _, err := os.Stat(binaryPath); err != nil {
		log.Info().Msg("skipping setup experience: swiftDialog is not installed")
		return nil
	}

	log.Info().Msg("checking setup experience status")

	// Poll the status endpoint. This also releases the device if we're done.
	payload, err := s.OrbitClient.GetSetupExperienceStatus(false)
	if err != nil {
		return err
	}

	// If swiftDialog isn't up yet, then launch it
	orgLogo := payload.OrgLogoURL
	if orgLogo == "" {
		orgLogo = "https://fleetdm.com/images/permanent/fleet-mark-color-40x40@4x.png"
	}

	if err := s.startSwiftDialogLegacy(binaryPath, orgLogo); err != nil {
		return err
	}

	// Defer this so that s.started is only false the first time this function runs.
	defer func() { s.started = true }()

	select {
	case <-s.closeChan:
		log.Info().Str("receiver", "setup_experiencer").Msg("swiftDialog closed")
		return nil
	default:
		// ok
	}

	// We're rendering the initial loading UI (shown while there are still profiles, bootstrap package,
	// and account configuration to verify) right off the bat, so we can just no-op if any of those
	// are not terminal

	log.Info().Msg("setup experience: checking for pending statuses")

	if payload.BootstrapPackage != nil {
		if payload.BootstrapPackage.Status != fleet.MDMBootstrapPackageFailed && payload.BootstrapPackage.Status != fleet.MDMBootstrapPackageInstalled {
			log.Info().Msg("setup experience: bootstrap package pending")
			return nil
		}
	}

	if isPending, name := anyProfilePending(payload.ConfigurationProfiles); isPending {
		log.Info().Msg(fmt.Sprintf("setup experience: profile pending: %s", name))
		return nil
	}

	if payload.AccountConfiguration != nil {
		if payload.AccountConfiguration.Status != fleet.MDMAppleStatusAcknowledged &&
			payload.AccountConfiguration.Status != fleet.MDMAppleStatusError &&
			payload.AccountConfiguration.Status != fleet.MDMAppleStatusCommandFormatError {

			log.Info().Msg("setup experience: account config pending")
			return nil
		}
	}

	// Note that we are setting this based on the current payload only just in case something
	// was removed from the payload that was there earlier(e.g. a deleted software title).
	allStepsDone := true

	// Log the payload JSON
	log.Debug().Msgf("setup experience: payload: %+v", payload)

	// Now render the UI for the software and script.
	if len(payload.Software) > 0 || payload.Script != nil {
		log.Info().Msg("setup experience: rendering software and script UI")

		var stepsDone int
		var prog uint
		var steps []*fleet.SetupExperienceStatusResult
		if len(payload.Software) > 0 {
			steps = payload.Software
		}

		if payload.Script != nil {
			steps = append(steps, payload.Script)
		}

		// Check for any items that were in the payload that are no longer there. This can happen
		// if a software title was deleted, for instance
		for uiStepName, uiStep := range s.uiSteps {
			uiStepExistsInPayload := false
			for _, step := range steps {
				if uiStep.Title == step.Name {
					uiStepExistsInPayload = true
					break
				}
			}
			if !uiStepExistsInPayload {
				log.Info().Msgf("Setup Experience: list item %s removed from payload", uiStep.Title)
				err = s.sd.DeleteListItemByTitle(uiStep.Title)
				if err != nil {
					log.Info().Err(err).Msg("deleting list item removed from payload from setup experience UI")
				}
				delete(s.uiSteps, uiStepName)
			}
		}

		for _, step := range steps {
			currentStepState := resultToListItem(step)
			if priorStepState, ok := s.uiSteps[step.Name]; ok {
				if currentStepState != priorStepState {
					// We only want to resend on change so we're not unnecessarily scrolling the UI
					err = s.sd.UpdateListItemByTitle(currentStepState.Title, currentStepState.StatusText, currentStepState.Status)
					if err != nil {
						log.Info().Err(err).Msg("updating list item in setup experience UI")
					}
				} else {
					log.Info().Msgf("setup experience: no change in status for %s", step.Name)
				}
			} else {
				err = s.sd.AddListItem(currentStepState)
				if err != nil {
					log.Info().Err(err).Msg("adding list item in setup experience UI")
				}
				s.uiSteps[step.Name] = currentStepState
			}

			if step.Status == fleet.SetupExperienceStatusFailure || step.Status == fleet.SetupExperienceStatusSuccess {
				stepsDone++
				// The swiftDialog progress bar is out of 100
				for range int(float32(1) / float32(len(steps)) * 100) {
					prog++
				}
			} else {
				allStepsDone = false
			}
		}

		if err = s.sd.UpdateProgress(prog); err != nil {
			log.Info().Err(err).Msg("updating progress bar in setup experience UI")
		}

		if err := s.sd.ShowList(); err != nil {
			log.Info().Err(err).Msg("showing progress bar in setup experience UI")
		}

		if err := s.sd.UpdateProgressText(fmt.Sprintf("%.0f%%", float32(stepsDone)/float32(len(steps))*100)); err != nil {
			log.Info().Err(err).Msg("updating progress text in setup experience UI")
		}

	}

	// If we get here, we can render the "done" UI.

	if allStepsDone {
		if err := s.sd.SetMessage(doneMessage); err != nil {
			log.Info().Err(err).Msg("setting message in setup experience UI")
		}

		if err := s.sd.CompleteProgress(); err != nil {
			log.Info().Err(err).Msg("completing progress bar in setup experience UI")
		}

		if len(payload.Software) > 0 || payload.Script != nil {
			// need to call this because SetMessage removes the list from the view for some reason :(
			if err := s.sd.ShowList(); err != nil {
				log.Info().Err(err).Msg("showing list in setup experience UI")
			}
		}

		if err := s.sd.UpdateProgressText("100%"); err != nil {
			log.Info().Err(err).Msg("updating progress text in setup experience UI")
		}

		if err := s.sd.EnableButton1(true); err != nil {
			log.Info().Err(err).Msg("enabling close button in setup experience UI")
		}

		// Sleep for a few seconds to let the user see the done message before closing
		// the UI
		time.Sleep(3 * time.Second)

		if err := s.sd.Quit(); err != nil {
			log.Info().Err(err).Msg("quitting setup experience UI on completion")
		}
	}

	return nil
}

func (s *SetupExperiencer) startSwiftDialogLegacy(binaryPath, orgLogo string) error {
	if s.started {
		log.Info().Msg("swiftDialog started")
		return nil
	}

	log.Info().Msg("creating swiftDialog instance")

	created := make(chan struct{})
	swiftDialog, err := swiftdialog.Create(context.Background(), binaryPath)
	if err != nil {
		return errors.New("creating swiftDialog instance: %w")
	}
	s.sd = swiftDialog

	iconSize, err := swiftdialog.GetIconSize(orgLogo)
	if err != nil {
		log.Error().Err(err).Msg("setup experience: getting icon size")
		iconSize = swiftdialog.DefaultIconSize
	}

	go func() {
		initOpts := &swiftdialog.SwiftDialogOptions{
			Title:            "none",
			Message:          "### Setting up your Mac...\n\nYour Mac is being configured by your organization using Fleet. This process may take some time to complete. Please don't attempt to restart or shut down the computer unless prompted to do so.",
			Icon:             orgLogo,
			MessageAlignment: swiftdialog.AlignmentCenter,
			CentreIcon:       true,
			Height:           "625",
			Big:              true,
			ProgressText:     "Configuring your device...",
			Button1Text:      "Close",
			Button1Disabled:  true,
			BlurScreen:       true,
			OnTop:            true,
			QuitKey:          "X", // Capital X to require command+shift+x
		}

		if err := s.sd.Start(context.Background(), initOpts, true); err != nil {
			log.Error().Err(err).Msg("starting swiftDialog instance")
		}

		if err = s.sd.ShowProgress(); err != nil {
			log.Error().Err(err).Msg("setting initial setup experience progress")
		}

		if err := s.sd.SetIconSize(iconSize); err != nil {
			log.Error().Err(err).Msg("setting initial setup experience icon size")
		}

		log.Debug().Msg("swiftDialog process started")
		created <- struct{}{}

		if _, err = s.sd.Wait(); err != nil {
			log.Error().Err(err).Msg("swiftdialog.Wait failed")
		}

		s.closeChan <- struct{}{}
	}()
	<-created
	return nil
}

func resultToListItem(result *fleet.SetupExperienceStatusResult) swiftdialog.ListItem {
	statusText := "Pending"
	status := swiftdialog.StatusWait

	switch result.Status {
	case fleet.SetupExperienceStatusFailure:
		status = swiftdialog.StatusFail
		statusText = "Failed"
	case fleet.SetupExperienceStatusSuccess:
		status = swiftdialog.StatusSuccess
		statusText = "Installed"
		if result.IsForScript() {
			statusText = "Ran"
		}
	}

	return swiftdialog.ListItem{
		Title:      result.Name,
		Status:     status,
		StatusText: statusText,
	}
}