sampler-fork/data/item.go

172 lines
3.4 KiB
Go

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
type Item struct {
Label string
SampleScript string
InitScript *string
TransformScript *string
Color *ui.Color
RateMs 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 {
return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err))
}
timeout := make(chan bool, 1)
go func() {
time.Sleep(time.Duration(i.RateMs))
timeout <- true
}()
var outputText strings.Builder
for {
select {
case output := <-i.InteractiveShell.Channel:
if !strings.Contains(output, i.SampleScript) && len(output) > 0 {
outputText.WriteString(output)
outputText.WriteString("\n")
}
case <-timeout:
sample := cleanupOutput(outputText.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 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
}