sampler-fork/data/int_pty.go

132 lines
2.4 KiB
Go
Raw Permalink Normal View History

2019-07-27 04:15:35 +00:00
//+build !windows
package data
import (
"bufio"
"fmt"
"github.com/kr/pty"
"github.com/lunixbochs/vtclean"
"io"
"os/exec"
"strings"
"time"
)
2019-07-27 04:15:35 +00:00
// PtyInteractiveShell represents PTY interactive shell sampling metadata
type PtyInteractiveShell struct {
item *Item
variables []string
cmd *exec.Cmd
file io.WriteCloser
ch chan string
errCount int
2019-06-21 04:32:02 +00:00
timeout time.Duration
}
func (s *PtyInteractiveShell) init() error {
2019-06-23 04:08:24 +00:00
cmd := exec.Command("sh", "-c", s.item.initScripts[0])
enrichEnvVariables(cmd, s.variables)
file, err := pty.Start(cmd)
if err != nil {
return err
}
scanner := bufio.NewScanner(file)
channel := make(chan string)
go func() {
for scanner.Scan() {
channel <- scanner.Text()
}
}()
s.cmd = cmd
s.file = file
s.ch = channel
_, err = file.Read(make([]byte, 4096))
if err != nil {
return err
}
time.Sleep(startupTimeout)
2019-06-23 04:08:24 +00:00
for i := 1; i < len(s.item.initScripts); i++ {
_, err = io.WriteString(s.file, fmt.Sprintf(" %s\n", s.item.initScripts[i]))
if err != nil {
return err
}
time.Sleep(startupTimeout) // TODO wait until cmd complete
}
return nil
}
func (s *PtyInteractiveShell) execute() (string, error) {
_, err := io.WriteString(s.file, fmt.Sprintf(" %s\n", s.item.sampleScript))
if err != nil {
s.errCount++
if s.errCount > errorThreshold {
_ = s.cmd.Wait()
_ = s.file.Close()
s.item.ptyShell = nil // restart session
}
2019-07-27 04:15:35 +00:00
return "", fmt.Errorf("failed to execute command: %s", err)
}
softTimeout := make(chan bool, 1)
hardTimeout := make(chan bool, 1)
go func() {
time.Sleep(s.getAwaitTimeout() / 2)
softTimeout <- true
time.Sleep(s.getAwaitTimeout() * 100)
hardTimeout <- true
}()
var builder strings.Builder
softTimeoutElapsed := false
await:
for {
select {
case out := <-s.ch:
cout := vtclean.Clean(out, false)
if len(cout) > 0 && !strings.Contains(cout, s.item.sampleScript) {
builder.WriteString(cout)
builder.WriteString("\n")
if softTimeoutElapsed {
break await
}
}
case <-softTimeout:
if builder.Len() > 0 {
break await
} else {
softTimeoutElapsed = true
}
case <-hardTimeout:
break await
}
}
sample := strings.TrimSpace(builder.String())
return s.item.transform(sample)
}
func (s *PtyInteractiveShell) getAwaitTimeout() time.Duration {
2019-06-21 04:32:02 +00:00
if s.timeout > maxAwaitTimeout {
return maxAwaitTimeout
2019-06-21 04:32:02 +00:00
} else if s.timeout < minAwaitTimeout {
return minAwaitTimeout
}
2019-06-21 04:32:02 +00:00
return s.timeout
}