diff --git a/src/sessionmanagerplugin/session/shellsession/shellsession.go b/src/sessionmanagerplugin/session/shellsession/shellsession.go index 991683c9..d7e7965c 100644 --- a/src/sessionmanagerplugin/session/shellsession/shellsession.go +++ b/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 +} diff --git a/src/sessionmanagerplugin/session/shellsession/shellsession_unix.go b/src/sessionmanagerplugin/session/shellsession/shellsession_unix.go index 783c2195..c99eee28 100644 --- a/src/sessionmanagerplugin/session/shellsession/shellsession_unix.go +++ b/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 diff --git a/src/sessionmanagerplugin/session/shellsession/shellsession_windows.go b/src/sessionmanagerplugin/session/shellsession/shellsession_windows.go index 71a3538c..ba70e9f0 100644 --- a/src/sessionmanagerplugin/session/shellsession/shellsession_windows.go +++ b/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