sampler-fork/data/item.go

172 lines
3.4 KiB
Go
Raw Normal View History

2019-01-31 23:40:05 +00:00
package data
import (
2019-04-07 15:09:24 +00:00
"bufio"
"errors"
"fmt"
2019-03-14 03:01:44 +00:00
ui "github.com/gizak/termui/v3"
"github.com/kr/pty"
"github.com/sqshq/sampler/config"
2019-04-07 15:09:24 +00:00
"io"
2019-03-21 02:23:08 +00:00
"os"
2019-01-31 23:40:05 +00:00
"os/exec"
"strings"
2019-04-07 15:09:24 +00:00
"time"
2019-01-31 23:40:05 +00:00
)
const interactiveShellStartupTimeout = 100 * time.Millisecond
2019-01-31 23:40:05 +00:00
type Item struct {
2019-04-07 15:09:24 +00:00
Label string
SampleScript string
InitScript *string
TransformScript *string
Color *ui.Color
2019-04-07 15:17:28 +00:00
RateMs int
2019-04-07 15:09:24 +00:00
InteractiveShell *InteractiveShell
}
2019-04-07 15:09:24 +00:00
type InteractiveShell struct {
Channel chan string
File io.WriteCloser
Cmd *exec.Cmd
2019-04-07 15:09:24 +00:00
}
2019-04-07 15:17:28 +00:00
func NewItems(cfgs []config.Item, rateMs int) []*Item {
2019-04-07 15:09:24 +00:00
items := make([]*Item, 0)
for _, i := range cfgs {
2019-04-07 15:09:24 +00:00
item := &Item{
Label: *i.Label,
SampleScript: *i.SampleScript,
InitScript: i.InitScript,
TransformScript: i.TransformScript,
2019-04-07 15:09:24 +00:00
Color: i.Color,
2019-04-07 15:17:28 +00:00
RateMs: rateMs,
2019-04-07 15:09:24 +00:00
}
items = append(items, item)
}
return items
2019-01-31 23:40:05 +00:00
}
2019-04-07 15:09:24 +00:00
func (i *Item) nextValue(variables []string) (string, error) {
2019-01-31 23:40:05 +00:00
2019-04-07 15:09:24 +00:00
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))
2019-04-07 15:09:24 +00:00
}
}
2019-03-21 02:23:08 +00:00
2019-04-07 15:09:24 +00:00
if i.InitScript != nil {
return i.executeInteractiveShellCmd(variables)
} else {
return i.executeCmd(variables, i.SampleScript)
2019-03-21 02:23:08 +00:00
}
2019-04-07 15:09:24 +00:00
}
func (i *Item) executeCmd(variables []string, script string) (string, error) {
2019-04-07 15:09:24 +00:00
cmd := exec.Command("sh", "-c", script)
2019-04-07 15:09:24 +00:00
enrichEnvVariables(cmd, variables)
2019-03-21 02:23:08 +00:00
output, err := cmd.Output()
2019-01-31 23:40:05 +00:00
if err != nil {
return "", err
}
return string(output), nil
2019-01-31 23:40:05 +00:00
}
2019-04-07 15:09:24 +00:00
func (i *Item) initInteractiveShell(variables []string) error {
cmd := exec.Command("sh", "-c", *i.InitScript)
enrichEnvVariables(cmd, variables)
file, err := pty.Start(cmd)
2019-04-07 15:09:24 +00:00
if err != nil {
return err
}
scanner := bufio.NewScanner(file)
channel := make(chan string)
2019-04-07 15:09:24 +00:00
go func() {
for scanner.Scan() {
channel <- scanner.Text()
2019-04-07 15:09:24 +00:00
}
}()
i.InteractiveShell = &InteractiveShell{
Channel: channel,
File: file,
Cmd: cmd,
2019-04-07 15:09:24 +00:00
}
_, err = file.Read(make([]byte, 4096))
2019-04-07 15:09:24 +00:00
if err != nil {
return err
}
time.Sleep(interactiveShellStartupTimeout)
2019-04-07 15:09:24 +00:00
return nil
}
func (i *Item) executeInteractiveShellCmd(variables []string) (string, error) {
_, err := io.WriteString(i.InteractiveShell.File, fmt.Sprintf(" %s\n", i.SampleScript))
2019-04-07 15:09:24 +00:00
if err != nil {
return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err))
2019-04-07 15:09:24 +00:00
}
timeout := make(chan bool, 1)
go func() {
time.Sleep(time.Duration(i.RateMs))
2019-04-07 15:09:24 +00:00
timeout <- true
}()
var outputText strings.Builder
2019-04-07 15:09:24 +00:00
for {
select {
case output := <-i.InteractiveShell.Channel:
if !strings.Contains(output, i.SampleScript) && len(output) > 0 {
outputText.WriteString(output)
outputText.WriteString("\n")
2019-04-07 15:09:24 +00:00
}
case <-timeout:
sample := cleanupOutput(outputText.String())
return i.transformInteractiveShellCmd(sample)
2019-04-07 15:09:24 +00:00
}
}
}
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
2019-04-07 15:09:24 +00:00
}
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
}