Add Escape Sequences to Shell Sessions

src/sessionmanagerplugin/session/shellsession/shellsession.go

@@ -40,6 +40,7 @@ type ShellSession struct {
 	// SizeData is used to store size data at session level to compare with new size.
 	SizeData          message.SizeData
 	originalSttyState bytes.Buffer
+	escapeTracking    ShellEscapeSequenceTracking
 }
 
 var GetTerminalSizeCall = func(fd int) (width int, height int, err error) {
@@ -62,6 +63,11 @@ func (s *ShellSession) Initialize(log log.T, sessionVar *session.Session) {
 		func(input []byte) {
 			s.DataChannel.OutputMessageHandler(log, s.Stop, s.SessionId, input)
 		})
+	s.escapeTracking = ShellEscapeSequenceTracking{
+		enabled: true,
+		newline: false,
+		escaped: false,
+	}
 }
 
 // StartSession takes input and write it to data channel
@@ -138,3 +144,83 @@ func (s ShellSession) ProcessStreamMessagePayload(log log.T, outputMessage messa
 	s.DisplayMode.DisplayMessage(log, outputMessage)
 	return true, nil
 }
+
+// Shell Session Escape Sequence Tracking Flags
+type ShellEscapeSequenceTracking struct {
+	enabled bool // whether the shell session should check for escape sequences
+	newline bool // whether the last character was a newline (the first half of the trigger)
+	escaped bool // whether the shell session is escaped (the second half of the trigger)
+}
+
+// Reset Escape Sequence due to finishing an escape sequence, or invalid escape sequence.
+func (s *ShellEscapeSequenceTracking) Reset() {
+	s.escaped = false
+	s.newline = false
+}
+
+// Disable checking for Escape Sequences for the rest of the connected Session.
+func (s *ShellEscapeSequenceTracking) Disable() {
+	s.enabled = false
+}
+
+// First half of trigger (newline) detected
+func (s *ShellEscapeSequenceTracking) HalfTrigger() {
+	s.newline = true
+}
+
+// Second half of trigger (~) detected
+func (s *ShellEscapeSequenceTracking) Trigger() {
+	if s.newline {
+		s.escaped = true
+	} else {
+		panic("Unexpected trigger, when prior newline missing")
+	}
+}
+
+// handleEscapeSequence process key presses looking for the escape sequence
+func (s *ShellSession) handleEscapeSequence(log log.T, stdinBytes []byte, stdinBytesLen int) (skipMessage bool, err error) {
+	const (
+		escape_help = "\nSupported escape sequence commands:\n" +
+			"~?  - this help message\n" +
+			"~~  - send the ~ character to the remote target\n" +
+			"~-  - disable escape sequences for the rest of this session\n" +
+			"~.  - disconnect and terminate session\n" +
+			"(Note that escapes are only recognized immediately after newline.)"
+	)
+
+	if s.escapeTracking.enabled {
+		if s.escapeTracking.newline && stdinBytesLen == 1 {
+			if s.escapeTracking.escaped {
+				switch stdinBytes[0] {
+				case '?': // help
+					println(escape_help)
+					s.escapeTracking.Reset()
+					return true, nil
+				case '.': // disconnect and terminate
+					if err := s.Session.TerminateSession(log); err != nil {
+						return true, err
+					}
+					return true, nil
+				case '-': // disable
+					s.escapeTracking.Disable()
+					return true, nil
+				case '~': // send explicit ~ character
+					s.escapeTracking.Reset()
+					return false, nil
+				}
+				s.escapeTracking.Reset()
+			} else if stdinBytes[0] == '~' {
+				s.escapeTracking.Trigger()
+				return true, nil
+			} else {
+				s.escapeTracking.Reset()
+			}
+		}
+
+		// If last sent bytes ends with newline, mark as possible escape sequence
+		if stdinBytes[stdinBytesLen-1] == '\n' || stdinBytes[stdinBytesLen-1] == '\r' {
+			s.escapeTracking.HalfTrigger()
+		}
+	}
+	return false, nil
+}

src/sessionmanagerplugin/session/shellsession/shellsession_unix.go

@@ -75,6 +75,13 @@ func (s *ShellSession) handleKeyboardInput(log log.T) (err error) {
 			break
 		}
 
+		if skip, err := s.handleEscapeSequence(log, stdinBytes, stdinBytesLen); err != nil {
+			log.Errorf("Escape sequence failure: %v", err)
+			s.Stop()
+		} else if skip {
+			continue
+		}
+
 		if err = s.Session.DataChannel.SendInputDataMessage(log, message.Output, stdinBytes[:stdinBytesLen]); err != nil {
 			log.Errorf("Failed to send UTF8 char: %v", err)
 			break

src/sessionmanagerplugin/session/shellsession/shellsession_windows.go

@@ -77,6 +77,12 @@ func (s *ShellSession) handleKeyboardInput(log log.T) (err error) {
 		}
 		if character != 0 {
 			charBytes := []byte(string(character))
+			if skip, err := s.handleEscapeSequence(log, charBytes, len(charBytes)); err != nil {
+				log.Errorf("Escape sequence failure: %v", err)
+				s.Stop()
+			} else if skip {
+				continue
+			}
 			if err = s.Session.DataChannel.SendInputDataMessage(log, message.Output, charBytes); err != nil {
 				log.Errorf("Failed to send UTF8 char: %v", err)
 				break
@@ -86,6 +92,12 @@ func (s *ShellSession) handleKeyboardInput(log log.T) (err error) {
 			if byteValue, ok := specialKeysInputMap[key]; ok {
 				keyBytes = byteValue
 			}
+			if skip, err := s.handleEscapeSequence(log, keyBytes, len(keyBytes)); err != nil {
+				log.Errorf("Escape sequence failure: %v", err)
+				s.Stop()
+			} else if skip {
+				continue
+			}
 			if err = s.Session.DataChannel.SendInputDataMessage(log, message.Output, keyBytes); err != nil {
 				log.Errorf("Failed to send UTF8 char: %v", err)
 				break