From 42683e6f4af77bc90373b47306e7b7f0418cbee9 Mon Sep 17 00:00:00 2001 From: sawka Date: Wed, 14 Sep 2022 23:10:35 -0700 Subject: [PATCH] working on remote pty output --- cmd/main-server.go | 30 +++++++++++ go.mod | 1 + go.sum | 2 + pkg/remote/remote.go | 108 ++++++++++++++++++++++++++-------------- pkg/sstore/updatebus.go | 5 +- 5 files changed, 108 insertions(+), 38 deletions(-) diff --git a/cmd/main-server.go b/cmd/main-server.go index 5add214a..fd89e92b 100644 --- a/cmd/main-server.go +++ b/cmd/main-server.go @@ -142,6 +142,35 @@ func HandleGetWindow(w http.ResponseWriter, r *http.Request) { return } +func HandleRemotePty(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin")) + w.Header().Set("Access-Control-Allow-Credentials", "true") + w.Header().Set("Vary", "Origin") + w.Header().Set("Cache-Control", "no-cache") + qvals := r.URL.Query() + remoteId := qvals.Get("remoteid") + if remoteId == "" { + w.WriteHeader(500) + w.Write([]byte(fmt.Sprintf("must specify remoteid"))) + return + } + if _, err := uuid.Parse(remoteId); err != nil { + w.WriteHeader(500) + w.Write([]byte(fmt.Sprintf("invalid remoteid: %v", err))) + return + } + realOffset, data, err := remote.ReadRemotePty(r.Context(), remoteId) + if err != nil { + w.WriteHeader(500) + w.Write([]byte(fmt.Sprintf("error reading ptyout file: %v", err))) + return + } + w.Header().Set("X-PtyDataOffset", strconv.FormatInt(realOffset, 10)) + w.WriteHeader(http.StatusOK) + w.Write(data) + return +} + func HandleGetPtyOut(w http.ResponseWriter, r *http.Request) { w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin")) w.Header().Set("Access-Control-Allow-Credentials", "true") @@ -333,6 +362,7 @@ func main() { go runWebSocketServer() gr := mux.NewRouter() gr.HandleFunc("/api/ptyout", HandleGetPtyOut) + gr.HandleFunc("/api/remote-pty", HandleRemotePty) gr.HandleFunc("/api/get-window", HandleGetWindow) gr.HandleFunc("/api/run-command", HandleRunCommand).Methods("GET", "POST", "OPTIONS") server := &http.Server{ diff --git a/go.mod b/go.mod index 3fe97fe9..d4a4489f 100644 --- a/go.mod +++ b/go.mod @@ -14,6 +14,7 @@ require ( require ( github.com/alessio/shellescape v1.4.1 // indirect + github.com/armon/circbuf v0.0.0-20190214190532-5111143e8da2 // indirect github.com/creack/pty v1.1.18 // indirect github.com/fsnotify/fsnotify v1.5.4 // indirect github.com/hashicorp/errwrap v1.1.0 // indirect diff --git a/go.sum b/go.sum index 1e458fc3..7465dfaa 100644 --- a/go.sum +++ b/go.sum @@ -119,6 +119,8 @@ github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kd github.com/apache/arrow/go/arrow v0.0.0-20210818145353-234c94e4ce64/go.mod h1:2qMFB56yOP3KzkB3PbYZ4AlUFg3a88F67TIx5lB/WwY= github.com/apache/arrow/go/arrow v0.0.0-20211013220434-5962184e7a30/go.mod h1:Q7yQnSMnLvcXlZ8RV+jwz/6y1rQTqbX6C82SndT52Zs= github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o= +github.com/armon/circbuf v0.0.0-20190214190532-5111143e8da2 h1:7Ip0wMmLHLRJdrloDxZfhMm0xrLXZS8+COSu2bXmEQs= +github.com/armon/circbuf v0.0.0-20190214190532-5111143e8da2/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o= github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8= github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY= github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8= diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go index 2527dffd..bff2b4df 100644 --- a/pkg/remote/remote.go +++ b/pkg/remote/remote.go @@ -17,6 +17,7 @@ import ( "syscall" "time" + "github.com/armon/circbuf" "github.com/creack/pty" "github.com/google/uuid" "github.com/scripthaus-dev/mshell/pkg/base" @@ -28,6 +29,7 @@ import ( const RemoteTypeMShell = "mshell" const DefaultTerm = "xterm-256color" const DefaultMaxPtySize = 1024 * 1024 +const CircBufSize = 64 * 1024 const MShellServerCommand = ` PATH=$PATH:~/.mshell; @@ -52,7 +54,6 @@ var GlobalStore *Store type Store struct { Lock *sync.Mutex Map map[string]*MShellProc // key=remoteid - Log *CircleLog CmdWaitMap map[base.CommandKey][]func() } @@ -66,6 +67,7 @@ type MShellProc struct { UName string Err error ControllingPty *os.File + PtyBuffer *circbuf.Buffer RunningCmds []base.CommandKey } @@ -85,21 +87,6 @@ type RemoteRuntimeState struct { RemoteIdx int64 `json:"remoteidx"` } -func logf(rem *sstore.RemoteType, fmtStr string, args ...interface{}) { - rname := rem.GetName() - str := fmt.Sprintf(fmtStr, args...) - fullStr := fmt.Sprintf("[remote %s] %s", rname, str) - fmt.Printf("%s\n", fullStr) - GlobalStore.Log.Add(fullStr) -} - -func logError(rem *sstore.RemoteType, err error) { - rname := rem.GetName() - fullStr := fmt.Sprintf("[remote %s] error: %v", rname, err) - fmt.Printf("%s\n", fullStr) - GlobalStore.Log.Add(fullStr) -} - func (state RemoteRuntimeState) IsConnected() bool { return state.Status == StatusConnected } @@ -129,7 +116,6 @@ func LoadRemotes(ctx context.Context) error { GlobalStore = &Store{ Lock: &sync.Mutex{}, Map: make(map[string]*MShellProc), - Log: MakeCircleLog(100), CmdWaitMap: make(map[base.CommandKey][]func()), } allRemotes, err := sstore.GetAllRemotes(ctx) @@ -168,6 +154,20 @@ func LoadRemoteById(ctx context.Context, remoteId string) error { return nil } +func ReadRemotePty(ctx context.Context, remoteId string) (int64, []byte, error) { + GlobalStore.Lock.Lock() + defer GlobalStore.Lock.Unlock() + msh := GlobalStore.Map[remoteId] + if msh == nil { + return 0, nil, nil + } + msh.Lock.Lock() + defer msh.Lock.Unlock() + barr := msh.PtyBuffer.Bytes() + offset := msh.PtyBuffer.TotalWritten() - int64(len(barr)) + return offset, barr, nil +} + func AddRemote(ctx context.Context, r *sstore.RemoteType) error { GlobalStore.Lock.Lock() defer GlobalStore.Lock.Unlock() @@ -423,7 +423,17 @@ func GetDefaultRemoteStateById(remoteId string) (*sstore.RemoteState, error) { } func MakeMShell(r *sstore.RemoteType) *MShellProc { - rtn := &MShellProc{Lock: &sync.Mutex{}, Remote: r, Status: StatusInit} + buf, err := circbuf.NewBuffer(CircBufSize) + if err != nil { + panic(err) // this should never happen (NewBuffer only returns an error if CirBufSize <= 0) + } + rtn := &MShellProc{ + Lock: &sync.Mutex{}, + Remote: r, + Status: StatusInit, + PtyBuffer: buf, + } + rtn.WriteToPtyBuffer("console for remote [%s]\n", r.GetName()) return rtn } @@ -492,21 +502,46 @@ func (msh *MShellProc) GetRemoteName() string { return msh.Remote.GetName() } +func (msh *MShellProc) WriteToPtyBuffer(strFmt string, args ...interface{}) { + msh.Lock.Lock() + defer msh.Lock.Unlock() + msh.writeToPtyBuffer_nolock(strFmt, args...) +} + +func (msh *MShellProc) writeToPtyBuffer_nolock(strFmt string, args ...interface{}) { + // inefficient string manipulation here and read of PtyBuffer, but these messages are rare, nbd + realStr := fmt.Sprintf(strFmt, args...) + if !strings.HasPrefix(realStr, "~") { + realStr = strings.ReplaceAll(realStr, "\n", "\r\n") + if !strings.HasSuffix(realStr, "\r\n") { + realStr = realStr + "\r\n" + } + realStr = "\033[0m\033[32m[scripthaus]\033[0m " + realStr + barr := msh.PtyBuffer.Bytes() + if len(barr) > 0 && barr[len(barr)-1] != '\n' { + realStr = "\r\n" + realStr + } + } else { + realStr = realStr[1:] + } + msh.PtyBuffer.Write([]byte(realStr)) +} + func (msh *MShellProc) Launch() { remoteCopy := msh.GetRemoteCopy() remoteName := remoteCopy.GetName() if remoteCopy.Archived { - logf(&remoteCopy, "cannot launch archived remote") + msh.WriteToPtyBuffer("cannot launch archived remote\n") return } - logf(&remoteCopy, "starting launch") + msh.WriteToPtyBuffer("connecting to %s\n", remoteCopy.GetName()) sshOpts := convertSSHOpts(remoteCopy.SSHOpts) sshOpts.SSHErrorsToTty = true ecmd := sshOpts.MakeSSHExecCmd(MShellServerCommand) cmdPty, err := msh.addControllingTty(ecmd) if err != nil { statusErr := fmt.Errorf("cannot attach controlling tty to mshell command: %w", err) - logError(&remoteCopy, statusErr) + msh.WriteToPtyBuffer("error, %s\n", statusErr.Error()) msh.setErrorStatus(statusErr) return } @@ -523,18 +558,19 @@ func (msh *MShellProc) Launch() { break } if readErr != nil { - logf(&remoteCopy, "error: reading from controlling-pty: %v", readErr) + msh.WriteToPtyBuffer("error reading from controlling-pty: %v\n", readErr) break } - readStr := string(buf[0:n]) - fmt.Printf("[c-pty %s] %d %q\n", remoteName, n, readStr) + msh.WithLock(func() { + msh.PtyBuffer.Write(buf[0:n]) + }) } }() if remoteName == "test2" { go func() { time.Sleep(2 * time.Second) cmdPty.Write([]byte(Test2Pw)) - fmt.Printf("[c-pty %s] wrote password!\n", remoteName) + msh.WriteToPtyBuffer("~[password sent]\r\n") }() } cproc, uname, err := shexec.MakeClientProc(ecmd) @@ -544,10 +580,10 @@ func (msh *MShellProc) Launch() { }) if err != nil { msh.setErrorStatus(err) - logf(&remoteCopy, "error connecting remote (%s): %v", msh.UName, err) + msh.WriteToPtyBuffer("error connecting to remote (uname=%q): %v\n", msh.UName, err) return } - fmt.Printf("connected remote %s\n", remoteCopy.GetName()) + msh.WriteToPtyBuffer("connected\n") msh.WithLock(func() { msh.ServerProc = cproc msh.Status = StatusConnected @@ -562,7 +598,7 @@ func (msh *MShellProc) Launch() { go msh.NotifyRemoteUpdate() } }) - logf(&remoteCopy, "remote disconnected exitcode=%d", exitCode) + msh.WriteToPtyBuffer("disconnected exitcode=%d\n", exitCode) }() go msh.ProcessPackets() return @@ -767,7 +803,7 @@ func (msh *MShellProc) notifyHangups_nolock() { func (msh *MShellProc) handleCmdDonePacket(donePk *packet.CmdDonePacketType) { update, err := sstore.UpdateCmdDonePk(context.Background(), donePk) if err != nil { - fmt.Printf("[error] updating cmddone: %v\n", err) + msh.WriteToPtyBuffer("[error] updating cmddone: %v\n", err) return } if update != nil { @@ -780,7 +816,7 @@ func (msh *MShellProc) handleCmdDonePacket(donePk *packet.CmdDonePacketType) { func (msh *MShellProc) handleCmdErrorPacket(errPk *packet.CmdErrorPacketType) { err := sstore.AppendCmdErrorPk(context.Background(), errPk) if err != nil { - fmt.Printf("cmderr> [remote %s] [error] adding cmderr: %v\n", msh.GetRemoteName(), err) + msh.WriteToPtyBuffer("cmderr> [remote %s] [error] adding cmderr: %v\n", msh.GetRemoteName(), err) return } return @@ -832,7 +868,7 @@ func (msh *MShellProc) ProcessPackets() { } err := sstore.HangupRunningCmdsByRemoteId(context.Background(), msh.Remote.RemoteId) if err != nil { - logf(msh.Remote, "calling HUP on cmds %v", err) + msh.writeToPtyBuffer_nolock("error calling HUP on cmds %v\n", err) } msh.notifyHangups_nolock() go msh.NotifyRemoteUpdate() @@ -851,7 +887,7 @@ func (msh *MShellProc) ProcessPackets() { } if pk.GetType() == packet.CmdDataPacketStr { dataPacket := pk.(*packet.CmdDataPacketType) - fmt.Printf("cmd-data> [remote %s] [%s] pty=%d run=%d\n", msh.GetRemoteName(), dataPacket.CK, dataPacket.PtyDataLen, dataPacket.RunDataLen) + msh.WriteToPtyBuffer("cmd-data> [remote %s] [%s] pty=%d run=%d\n", msh.GetRemoteName(), dataPacket.CK, dataPacket.PtyDataLen, dataPacket.RunDataLen) continue } if pk.GetType() == packet.CmdDonePacketStr { @@ -865,20 +901,20 @@ func (msh *MShellProc) ProcessPackets() { } if pk.GetType() == packet.MessagePacketStr { msgPacket := pk.(*packet.MessagePacketType) - fmt.Printf("# [remote %s] [%s] %s\n", msh.GetRemoteName(), msgPacket.CK, msgPacket.Message) + msh.WriteToPtyBuffer("msg> [remote %s] [%s] %s\n", msh.GetRemoteName(), msgPacket.CK, msgPacket.Message) continue } if pk.GetType() == packet.RawPacketStr { rawPacket := pk.(*packet.RawPacketType) - fmt.Printf("stderr> [remote %s] %s\n", msh.GetRemoteName(), rawPacket.Data) + msh.WriteToPtyBuffer("stderr> [remote %s] %s\n", msh.GetRemoteName(), rawPacket.Data) continue } if pk.GetType() == packet.CmdStartPacketStr { startPk := pk.(*packet.CmdStartPacketType) - fmt.Printf("start> [remote %s] reqid=%s (%p)\n", msh.GetRemoteName(), startPk.RespId, msh.ServerProc.Output) + msh.WriteToPtyBuffer("start> [remote %s] reqid=%s (%p)\n", msh.GetRemoteName(), startPk.RespId, msh.ServerProc.Output) continue } - fmt.Printf("MSH> [remote %s] unhandled packet %s\n", msh.GetRemoteName(), packet.AsString(pk)) + msh.WriteToPtyBuffer("MSH> [remote %s] unhandled packet %s\n", msh.GetRemoteName(), packet.AsString(pk)) } } diff --git a/pkg/sstore/updatebus.go b/pkg/sstore/updatebus.go index 69dece4d..524b5af3 100644 --- a/pkg/sstore/updatebus.go +++ b/pkg/sstore/updatebus.go @@ -15,8 +15,9 @@ type UpdatePacket interface { } type PtyDataUpdate struct { - SessionId string `json:"sessionid"` - CmdId string `json:"cmdid"` + SessionId string `json:"sessionid,omitempty"` + CmdId string `json:"cmdid,omitempty"` + RemoteId string `json:"remoteid,omitempty"` PtyPos int64 `json:"ptypos"` PtyData64 string `json:"ptydata64"` PtyDataLen int64 `json:"ptydatalen"`