package data import ( "bufio" "errors" "fmt" ui "github.com/gizak/termui/v3" "github.com/kr/pty" "github.com/sqshq/sampler/config" "io" "os" "os/exec" "strings" "time" ) const ( interactiveShellStartupTimeout = 100 * time.Millisecond interactiveShellMinAwaitTimeout = 100 * time.Millisecond interactiveShellMaxAwaitTimeout = 1 * time.Second interactiveShellErrorThreshold = 10 ) type Item struct { label string sampleScript string initScript *string transformScript *string color *ui.Color rateMs int errorsCount int interactiveShell *InteractiveShell } type InteractiveShell struct { Channel chan string File io.WriteCloser Cmd *exec.Cmd } func NewItems(cfgs []config.Item, rateMs int) []*Item { items := make([]*Item, 0) for _, i := range cfgs { item := &Item{ label: *i.Label, sampleScript: *i.SampleScript, initScript: i.InitScript, transformScript: i.TransformScript, color: i.Color, rateMs: rateMs, } items = append(items, item) } return items } func (i *Item) nextValue(variables []string) (string, error) { if i.initScript != nil && i.interactiveShell == nil { err := i.initInteractiveShell(variables) if err != nil { return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err)) } } if i.initScript != nil { return i.executeInteractiveShellCmd(variables) } else { return i.executeCmd(variables, i.sampleScript) } } func (i *Item) executeCmd(variables []string, script string) (string, error) { cmd := exec.Command("sh", "-c", script) enrichEnvVariables(cmd, variables) output, err := cmd.Output() if err != nil { return "", err } return string(output), nil } func (i *Item) initInteractiveShell(variables []string) error { cmd := exec.Command("sh", "-c", *i.initScript) enrichEnvVariables(cmd, 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() } }() i.interactiveShell = &InteractiveShell{ Channel: channel, File: file, Cmd: cmd, } _, err = file.Read(make([]byte, 4096)) if err != nil { return err } time.Sleep(interactiveShellStartupTimeout) return nil } func (i *Item) executeInteractiveShellCmd(variables []string) (string, error) { _, err := io.WriteString(i.interactiveShell.File, fmt.Sprintf(" %s\n", i.sampleScript)) if err != nil { i.errorsCount++ if i.errorsCount > interactiveShellErrorThreshold { i.interactiveShell = nil // restart session i.errorsCount = 0 } return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err)) } softTimeout := make(chan bool, 1) hardTimeout := make(chan bool, 1) go func() { time.Sleep(i.getAwaitTimeout() / 4) softTimeout <- true time.Sleep(i.getAwaitTimeout() * 100) hardTimeout <- true }() var builder strings.Builder softTimeoutElapsed := false await: for { select { case output := <-i.interactiveShell.Channel: o := cleanupOutput(output) if len(o) > 0 && !strings.Contains(o, i.sampleScript) { builder.WriteString(o) 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 i.transformInteractiveShellCmd(sample) } func (i *Item) transformInteractiveShellCmd(sample string) (string, error) { if i.transformScript != nil && len(sample) > 0 { return i.executeCmd([]string{"sample=" + sample}, *i.transformScript) } return sample, nil } func (i *Item) getAwaitTimeout() time.Duration { timeout := time.Duration(i.rateMs) * time.Millisecond if timeout > interactiveShellMaxAwaitTimeout { return interactiveShellMaxAwaitTimeout } else if timeout < interactiveShellMinAwaitTimeout { return interactiveShellMinAwaitTimeout } return timeout } func enrichEnvVariables(cmd *exec.Cmd, variables []string) { cmd.Env = os.Environ() for _, variable := range variables { cmd.Env = append(cmd.Env, variable) } } func cleanupOutput(output string) string { s := strings.TrimSpace(output) if idx := strings.Index(s, "\r"); idx != -1 { return s[idx+1:] } return s }