session.go

Documentation: edge-infra.dev/pkg/sds/emergencyaccess/emulator

     1  package emulator
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"strings"
     8  	"time"
     9  
    10  	"github.com/c-bata/go-prompt"
    11  	"github.com/google/shlex"
    12  
    13  	"edge-infra.dev/pkg/lib/fog"
    14  	"edge-infra.dev/pkg/sds/emergencyaccess/msgdata"
    15  	"edge-infra.dev/pkg/sds/lib/colors"
    16  )
    17  
    18  // This file contains all the private fields and methods for the session prompt.
    19  // the session prompt is responsible for:
    20  //   - handling and passing commands during a connected session
    21  //   - pausing display of data on the emulator when sp is active
    22  //   - returning feedback to the user incase of a bad command
    23  var (
    24  	internalErrChan = make(chan error)
    25  	timeoutChan     = make(chan bool)
    26  	promptActiveC   = make(chan bool)
    27  	livePrefixState struct {
    28  		livePrefix string
    29  		isEnabled  bool
    30  	}
    31  )
    32  
    33  const (
    34  	activeSessionPrefix   = ">>> "
    35  	darkModeSessionPrefix = "ddd> "
    36  )
    37  
    38  func changeLivePrefix() (string, bool) {
    39  	return livePrefixState.livePrefix, livePrefixState.isEnabled
    40  }
    41  
    42  func sessionPromptCompleter(in prompt.Document) []prompt.Suggest {
    43  	s := []prompt.Suggest{
    44  		{Text: "help", Description: "Display help"},
    45  		{Text: "exit", Description: "Return to Remote CLI prompt"},
    46  	}
    47  
    48  	// this means the session is active and we can show prompt suggestions
    49  	if livePrefixState.livePrefix == activeSessionPrefix || livePrefixState.livePrefix == darkModeSessionPrefix {
    50  		return prompt.FilterHasPrefix(s, in.CurrentLineBeforeCursor(), true)
    51  	}
    52  	// by "default" don't show anything
    53  	return []prompt.Suggest{}
    54  }
    55  
    56  // Used to split the command string from emulator command options
    57  // unit separator ascii 1f is unlikely to be used when typing terminal commands, although it is not guaranteed
    58  const sendOptionsSplitter = string('\u001f') + "||||" + string('\u001f')
    59  
    60  // Contains options used on the emulator side when dealing with commands and responses
    61  type CommandOpts struct {
    62  	// Specifies which file the output message should be sent to
    63  	FileRedirection string
    64  }
    65  
    66  // launches the display and input routines to display data asynchronously.
    67  func (em *Emulator) startSession() {
    68  	fmt.Println(colors.Text("Session prompt.'Ctrl-D', 'end', 'exit' or 'q' to exit. ", colors.BgGreen))
    69  	//passing common concolors.Text to both functions
    70  	ctx := em.cls.GetSessionContext()
    71  
    72  	go func() {
    73  		em.unPauseSessionPrompt()
    74  		em.sessionPrompt().Run()
    75  
    76  		// end the session and context
    77  		err := em.cls.End()
    78  		if err != nil {
    79  			em.log.Error(err, "ending session")
    80  		}
    81  	}()
    82  
    83  	go em.display(ctx)
    84  	em.handleExitEvents(ctx)
    85  }
    86  
    87  // creates a session prompt
    88  func (em *Emulator) sessionPrompt() *prompt.Prompt {
    89  	opts := []prompt.Option{
    90  		// utility
    91  		prompt.OptionHistory(em.sessionHistory.history),
    92  		prompt.OptionSetExitCheckerOnInput(em.sessionPromptExitHandler),
    93  		prompt.OptionLivePrefix(changeLivePrefix),
    94  
    95  		// keybinding
    96  		prompt.OptionAddKeyBind(append(sessionKeyBindings, commonKeyBindings...)...),
    97  
    98  		// colors
    99  		prompt.OptionPreviewSuggestionTextColor(prompt.Blue),
   100  		prompt.OptionSelectedSuggestionBGColor(prompt.LightGray),
   101  		prompt.OptionSuggestionBGColor(prompt.DarkGray),
   102  	}
   103  	darkmode := getDarkmode(em.cls)
   104  	if darkmode {
   105  		opts = append(opts,
   106  			prompt.OptionPrefixTextColor(prompt.Fuchsia),
   107  			prompt.OptionPrefix(darkModeSessionPrefix),
   108  		)
   109  	} else {
   110  		opts = append(opts,
   111  			prompt.OptionPrefixTextColor(prompt.Yellow),
   112  			prompt.OptionPrefix(activeSessionPrefix),
   113  		)
   114  	}
   115  	return prompt.New(em.sessionPromptExecutor,
   116  		sessionPromptCompleter,
   117  		opts...,
   118  	)
   119  }
   120  
   121  // similar to connect prompt executor.
   122  func (em *Emulator) sessionPromptExecutor(in string) {
   123  	if em.sessionPromptPaused {
   124  		// this means the session prompt is inactive
   125  		em.unPauseSessionPrompt()
   126  		promptActiveC <- true // pauses the display channel
   127  		return
   128  	}
   129  
   130  	if in == "help" {
   131  		fmt.Println(sessionHelp)
   132  		return
   133  	}
   134  	if em.handleBreakLineExit(in, true) {
   135  		return
   136  	}
   137  	in = strings.TrimRight(in, "\r")
   138  
   139  	err := em.sessionHistory.updateHistory(in, historyFileLimit)
   140  	if err != nil {
   141  		em.log.Error(err, "Error updating session history file")
   142  	}
   143  
   144  	command, opts, err := processSessionInput(in)
   145  	if err != nil {
   146  		fmt.Println(colors.Text("Error processing command: %s", colors.BgRed, in))
   147  		em.displaySessionError(err)
   148  		return
   149  	}
   150  
   151  	commandID, err := em.cls.Send(command)
   152  	if err != nil {
   153  		fmt.Println(colors.Text("Error with last submitted command: %q", colors.BgRed, command))
   154  		em.displaySessionError(err)
   155  	}
   156  
   157  	// Possibly shouldn't save if there is an error?
   158  	em.commandOptions[commandID] = opts
   159  
   160  	// deactivate the session prompt by default
   161  	em.pauseSessionPrompt()
   162  	promptActiveC <- false // tells the display channel it can display again
   163  }
   164  
   165  // Type implements UserError and error, printing the same string for both
   166  type userErr string
   167  
   168  func (err userErr) Error() string {
   169  	return string(err)
   170  }
   171  
   172  func (err userErr) UserError() []string {
   173  	return []string{string(err)}
   174  }
   175  
   176  func processSessionInput(input string) (string, CommandOpts, error) {
   177  	vals := strings.Split(input, sendOptionsSplitter)
   178  	if len(vals) < 2 {
   179  		// len vals must be at least 1 if sendOptionsSplitter is not empty,
   180  		// so safe to access first element
   181  		return vals[0], CommandOpts{}, nil
   182  	}
   183  
   184  	if len(vals) > 2 {
   185  		return "", CommandOpts{}, userErr("Unexpected number of command options")
   186  	}
   187  
   188  	opts, err := processCommandOptions(vals[1])
   189  	return vals[0], opts, err
   190  }
   191  
   192  func processCommandOptions(options string) (CommandOpts, error) {
   193  	vals, err := shlex.Split(options)
   194  	if err != nil {
   195  		return CommandOpts{}, err
   196  	}
   197  
   198  	if len(vals) == 0 {
   199  		return CommandOpts{}, userErr("No command options specified")
   200  	}
   201  
   202  	switch vals[0] {
   203  	case ">":
   204  		if len(vals) != 2 {
   205  			return CommandOpts{}, userErr("Unexpected number of arguments for File Redirection argument: " + options)
   206  		}
   207  		return CommandOpts{FileRedirection: vals[1]}, nil
   208  	default:
   209  		return CommandOpts{}, userErr(fmt.Sprintf("Unknown option specifier: %s", vals[0]))
   210  	}
   211  }
   212  
   213  func (em *Emulator) sessionPromptExitHandler(in string, breakline bool) bool {
   214  	// exit commands
   215  	select {
   216  	default:
   217  
   218  		return em.handleBreakLineExit(in, breakline)
   219  	// exit the session prompt without breaking the connect prompt.
   220  	case <-internalErrChan:
   221  		em.pauseSessionPrompt()
   222  		return true
   223  	case <-timeoutChan:
   224  		em.pauseSessionPrompt()
   225  		return true
   226  	}
   227  }
   228  
   229  func (em Emulator) handleBreakLineExit(in string, breakline bool) bool {
   230  	if breakline {
   231  		for _, item := range exitCommands {
   232  			if in == item {
   233  				em.pauseSessionPrompt()
   234  				return true
   235  			}
   236  		}
   237  	}
   238  	return false
   239  }
   240  
   241  // Blocking function that waits for the session context done channel to close or
   242  // for a session timeout
   243  func (em Emulator) handleExitEvents(ctx context.Context) {
   244  	ticker := time.NewTicker(1 * time.Second)
   245  	log := fog.FromContext(ctx)
   246  	for {
   247  		select {
   248  		case <-ctx.Done():
   249  			log.Info("handle exit events exiting normally")
   250  			return
   251  		case <-ticker.C:
   252  			// periodically check the idle time
   253  			if em.cls.IdleTime() > em.config.sessionTimeout {
   254  				// display the timeout notice from main routine
   255  				displayTimeoutNotice()
   256  
   257  				// wait for session prompt to return
   258  				timeoutChan <- true
   259  				return
   260  			}
   261  		}
   262  	}
   263  }
   264  
   265  func displayExitError() {
   266  	fmt.Println(colors.Text("\rSomething went wrong with the connection. See the remotecli.log log file in the current directory for further details", colors.BgRed))
   267  	fmt.Println(colors.Text("Press any key to return to Remote CLI prompt to attempt reconnection.", colors.FgRed))
   268  }
   269  
   270  func displayTimeoutNotice() {
   271  	fmt.Println(colors.Text("\rConnection timed out.", colors.BgRed))
   272  	fmt.Println(colors.Text("Press any key to return to Remote CLI prompt.", colors.FgRed))
   273  }
   274  
   275  // Blocking function that loops over the display channel and checks for returned responses
   276  // Exits when the context is done or the display channel is closed
   277  func (em Emulator) display(ctx context.Context) {
   278  	for {
   279  		select {
   280  		case response, ok := <-em.cls.GetDisplayChannel():
   281  			if !ok {
   282  				// Likely (but not guaranteed) this path will only be encountered
   283  				// when there is a subscription error, as the context done should
   284  				// be closed before the display channel
   285  				displayExitError()
   286  
   287  				// wait for session prompt exit handler to read this
   288  				internalErrChan <- nil
   289  
   290  				// return once the session prompt has exited
   291  				return
   292  			}
   293  			em.handleResponse(response)
   294  		case pause := <-promptActiveC:
   295  			if pause {
   296  				<-promptActiveC
   297  			}
   298  		case <-ctx.Done():
   299  			return
   300  		}
   301  	}
   302  }
   303  
   304  // Function handles all response messages and displays them on the console
   305  func (em *Emulator) handleResponse(response msgdata.CommandResponse) {
   306  	if response.Data().Type != "OUTPUT" {
   307  		return
   308  	}
   309  
   310  	display := true
   311  	if opts, ok := em.commandOptions[response.Attributes().ReqMsgID]; ok {
   312  		display = em.handleResponseOptions(response, opts)
   313  		delete(em.commandOptions, response.Attributes().ReqMsgID)
   314  	}
   315  
   316  	if display {
   317  		fmt.Println(response.Data().Output)
   318  	}
   319  
   320  	if response.Data().ExitCode == 0 {
   321  		fmt.Println(colors.BufferedText("Exit code: %d", colors.BgGreen, response.Data().ExitCode))
   322  	} else {
   323  		fmt.Println(colors.BufferedText("Exit code: %d", colors.BgRed, response.Data().ExitCode))
   324  	}
   325  }
   326  
   327  // Function executes all options for a particular response, returns boolean
   328  // indicating if the full response body should be displayed or only a summary
   329  func (em *Emulator) handleResponseOptions(response msgdata.CommandResponse, opts CommandOpts) bool {
   330  	display := true
   331  	if opts.FileRedirection != "" {
   332  		em.handleResponseFileRedirection(response, opts.FileRedirection)
   333  		display = false
   334  	}
   335  
   336  	return display
   337  }
   338  
   339  func (em *Emulator) handleResponseFileRedirection(response msgdata.CommandResponse, file string) {
   340  	err := os.WriteFile(file, []byte(response.Data().Output), 0644)
   341  	if err != nil {
   342  		em.log.Error(err, "failed writing response to file", "commandID", response.Attributes().ReqMsgID)
   343  		fmt.Println(colors.Text("Error occurred while saving response to file. See logs for details", colors.BgRed))
   344  	}
   345  }
   346  
   347  func (em *Emulator) pauseSessionPrompt() {
   348  	livePrefixState.isEnabled = true
   349  	livePrefixState.livePrefix = ""
   350  	em.sessionPromptPaused = true
   351  }
   352  
   353  func (em *Emulator) unPauseSessionPrompt() {
   354  	livePrefixState.isEnabled = true
   355  	livePrefixState.livePrefix = activeSessionPrefix
   356  	if getDarkmode(em.cls) {
   357  		livePrefixState.livePrefix = darkModeSessionPrefix
   358  	}
   359  	em.sessionPromptPaused = false
   360  }
   361
View as plain text