/
ssh_shell_conn.go
executable file
·191 lines (172 loc) · 5.01 KB
/
ssh_shell_conn.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
package ssh
import (
"bytes"
"github.com/gorilla/websocket"
"github.com/sirupsen/logrus"
"golang.org/x/crypto/ssh"
"io"
"sync"
"time"
)
// copy data from WebSocket to ssh server
// and copy data from ssh server to WebSocket
// write data to WebSocket
// the data comes from ssh server.
type wsBufferWriter struct {
buffer bytes.Buffer
mu sync.Mutex
}
// implement Write interface to write bytes from ssh server into bytes.Buffer.
func (w *wsBufferWriter) Write(p []byte) (int, error) {
w.mu.Lock()
defer w.mu.Unlock()
return w.buffer.Write(p)
}
const (
wsMsgCmd = "cmd"
wsMsgResize = "resize"
)
type wsMsg struct {
Type string `json:"type"`
Cmd string `json:"cmd"`
Cols int `json:"cols"`
Rows int `json:"rows"`
}
// connect to ssh server using ssh session.
type SshConn struct {
// calling Write() to write data into ssh server
StdinPipe io.WriteCloser
// Write() be called to receive data from ssh server
ComboOutput *wsBufferWriter
Session *ssh.Session
}
//flushComboOutput flush ssh.session combine output into websocket response
func flushComboOutput(w *wsBufferWriter, wsConn *websocket.Conn) error {
if w.buffer.Len() != 0 {
err := wsConn.WriteMessage(websocket.TextMessage, w.buffer.Bytes())
if err != nil {
return err
}
w.buffer.Reset()
}
return nil
}
// setup ssh shell session
// set Session and StdinPipe here,
// and the Session.Stdout and Session.Sdterr are also set.
func NewSshConn(cols, rows int, sshClient *ssh.Client) (*SshConn, error) {
sshSession, err := sshClient.NewSession()
if err != nil {
return nil, err
}
// we set stdin, then we can write data to ssh server via this stdin.
// but, as for reading data from ssh server, we can set Session.Stdout and Session.Stderr
// to receive data from ssh server, and write back to somewhere.
stdinP, err := sshSession.StdinPipe()
if err != nil {
return nil, err
}
comboWriter := new(wsBufferWriter)
//ssh.stdout and stderr will write output into comboWriter
sshSession.Stdout = comboWriter
sshSession.Stderr = comboWriter
modes := ssh.TerminalModes{
ssh.ECHO: 1, // disable echo
ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
}
// Request pseudo terminal
if err := sshSession.RequestPty("xterm", rows, cols, modes); err != nil {
return nil, err
}
// Start remote shell
if err := sshSession.Shell(); err != nil {
return nil, err
}
return &SshConn{StdinPipe: stdinP, ComboOutput: comboWriter, Session: sshSession}, nil
}
func (s *SshConn) Close() {
if s.Session != nil {
s.Session.Close()
}
}
//ReceiveWsMsg receive websocket msg do some handling then write into ssh.session.stdin
func (ssConn *SshConn) ReceiveWsMsg(wsConn *websocket.Conn, logBuff *bytes.Buffer, exitCh chan bool) {
//unmashal bytes into struct
msgObj := wsMsg{
Type: "cmd",
Cmd: "",
Rows: 50,
Cols: 180,
}
//tells other go routine quit
defer setQuit(exitCh)
for {
select {
case <-exitCh:
return
default:
//read websocket msg
_, wsData, err := wsConn.ReadMessage()
if err != nil {
logrus.WithError(err).Error("reading webSocket message failed")
return
}
//if err := json.Unmarshal(wsData, &msgObj); err != nil {
// logrus.WithError(err).WithField("wsData", string(wsData)).Error("unmarshal websocket message failed")
//}
switch msgObj.Type {
case wsMsgResize:
//handle xterm.js size change
if msgObj.Cols > 0 && msgObj.Rows > 0 {
if err := ssConn.Session.WindowChange(msgObj.Rows, msgObj.Cols); err != nil {
logrus.WithError(err).Error("ssh pty change windows size failed")
}
}
case wsMsgCmd:
//handle xterm.js stdin
//decodeBytes, err := base64.StdEncoding.DecodeString(msgObj.Cmd)
decodeBytes := wsData
if err != nil {
logrus.WithError(err).Error("websock cmd string base64 decoding failed")
}
if _, err := ssConn.StdinPipe.Write(decodeBytes); err != nil {
logrus.WithError(err).Error("ws cmd bytes write to ssh.stdin pipe failed")
}
//write input cmd to log buffer
if _, err := logBuff.Write(decodeBytes); err != nil {
logrus.WithError(err).Error("write received cmd into log buffer failed")
}
}
}
}
}
func (ssConn *SshConn) SendComboOutput(wsConn *websocket.Conn, exitCh chan bool) {
//tells other go routine quit
defer setQuit(exitCh)
//every 120ms write combine output bytes into websocket response
tick := time.NewTicker(time.Millisecond * time.Duration(120))
//for range time.Tick(120 * time.Millisecond){}
defer tick.Stop()
for {
select {
case <-tick.C:
//write combine output bytes into websocket response
if err := flushComboOutput(ssConn.ComboOutput, wsConn); err != nil {
logrus.WithError(err).Error("ssh sending combo output to webSocket failed")
return
}
case <-exitCh:
return
}
}
}
func (ssConn *SshConn) SessionWait(quitChan chan bool) {
if err := ssConn.Session.Wait(); err != nil {
logrus.WithError(err).Error("ssh session wait failed")
setQuit(quitChan)
}
}
func setQuit(ch chan bool) {
ch <- true
}