-
Notifications
You must be signed in to change notification settings - Fork 0
/
shell_finder.go
74 lines (68 loc) · 1.7 KB
/
shell_finder.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
package main
import (
"io"
"strings"
"github.com/acarl005/stripansi"
)
// ShellFindReader looks for the typical shell idicator (#,$,>)
// to appear at the end of stdout, then signalling by closing
// Found channel.
type ShellFindReader struct {
reader io.Reader
builder *strings.Builder
teeReader io.Reader
Found chan struct{}
SkippedEcho chan struct{}
Enabled bool
}
func NewShellFindReader(r io.Reader) *ShellFindReader {
builder := &strings.Builder{}
return &ShellFindReader{
reader: r,
builder: builder,
teeReader: io.TeeReader(r, builder),
Found: make(chan struct{}),
SkippedEcho: make(chan struct{}),
Enabled: true,
}
}
func (sfr *ShellFindReader) Read(p []byte) (int, error) {
select {
case <-sfr.Found:
if !sfr.Enabled {
discard := make([]byte, len(p))
_, err := sfr.teeReader.Read(discard)
accumulated := strings.TrimSpace(
sfr.builder.String(),
)
accumulated = stripansi.Strip(accumulated)
if accumulated != "" {
if strings.Contains(accumulated, mittenMarker) {
sfr.Enabled = true
sfr.builder.Reset()
close(sfr.SkippedEcho)
}
}
return 0, err
}
return sfr.reader.Read(p)
default:
n, err := sfr.teeReader.Read(p)
accumulated := sfr.builder.String()
accumulated = stripansi.Strip(accumulated)
accumulated = strings.TrimSpace(accumulated)
if accumulated != "" {
switch rune(accumulated[len(accumulated)-1]) {
case '$', '#', '>', '%':
// Skip the two last characters are the same (banners)
if accumulated[len(accumulated)-1] == accumulated[len(accumulated)-2] {
break
}
close(sfr.Found)
sfr.Enabled = false
sfr.builder.Reset()
}
}
return n, err
}
}