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
|
package handlers
import (
"bytes"
"encoding/base64"
"fmt"
"io"
"strings"
"time"
"github.com/mimecast/dtail/internal"
"github.com/mimecast/dtail/internal/io/dlog"
"github.com/mimecast/dtail/internal/protocol"
)
type baseHandler struct {
done *internal.Done
server string
shellStarted bool
commands chan string
receiveBuf bytes.Buffer
status int
}
func (h *baseHandler) String() string {
return fmt.Sprintf("baseHandler(%s,server:%s,shellStarted:%v,status:%d)@%p",
h.done,
h.server,
h.shellStarted,
h.status,
h,
)
}
func (h *baseHandler) Server() string {
return h.server
}
func (h *baseHandler) Status() int {
return h.status
}
// SendMessage to the server.
func (h *baseHandler) SendMessage(command string) error {
encoded := base64.StdEncoding.EncodeToString([]byte(command))
dlog.Client.Debug("Sending command", h.server, command, encoded)
select {
case h.commands <- fmt.Sprintf("protocol %s base64 %v;", protocol.ProtocolCompat, encoded):
case <-time.After(time.Second * 5):
return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded)
case <-h.Done():
return nil
}
return nil
}
// Read data from the dtail server via Writer interface.
func (h *baseHandler) Write(p []byte) (n int, err error) {
for _, b := range p {
switch b {
case protocol.MessageDelimiter:
message := h.receiveBuf.String()
h.handleMessage(message)
h.receiveBuf.Reset()
default:
h.receiveBuf.WriteByte(b)
}
}
return len(p), nil
}
// Send data to the dtail server via Reader interface.
func (h *baseHandler) Read(p []byte) (n int, err error) {
select {
case command := <-h.commands:
n = copy(p, []byte(command))
case <-h.Done():
return 0, io.EOF
}
return
}
func (h *baseHandler) handleMessage(message string) {
if len(message) > 0 && message[0] == '.' {
h.handleHiddenMessage(message)
return
}
dlog.Client.Raw(message)
}
// Handle messages received from server which are not meant to be displayed
// to the end user.
func (h *baseHandler) handleHiddenMessage(message string) {
switch {
case strings.HasPrefix(message, ".syn close connection"):
go h.SendMessage(".ack close connection")
h.Shutdown()
}
}
func (h *baseHandler) Done() <-chan struct{} {
return h.done.Done()
}
func (h *baseHandler) Shutdown() {
h.done.Shutdown()
}
|
