sampler-fork/data/item.go

183 lines
3.5 KiB
Go

package data
import (
"bufio"
"errors"
ui "github.com/gizak/termui/v3"
"github.com/sqshq/sampler/config"
"io"
"os"
"os/exec"
"strings"
"time"
)
type Item struct {
Label string
SampleScript string
InitScript *string
TransformScript *string
Color *ui.Color
RateMs int
InteractiveShell *InteractiveShell
}
type InteractiveShell struct {
StdoutCh chan string
StderrCh chan string
Stdin 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 "", 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)
stdout, err := cmd.StdoutPipe()
if err != nil {
return err
}
stderr, err := cmd.StderrPipe()
if err != nil {
return err
}
stdin, err := cmd.StdinPipe()
if err != nil {
return err
}
stdoutScanner := bufio.NewScanner(stdout)
stderrScanner := bufio.NewScanner(stderr)
stdoutCh := make(chan string)
stderrCh := make(chan string)
go func() {
for stdoutScanner.Scan() {
stdoutCh <- stdoutScanner.Text()
stderrCh <- stderrScanner.Text()
}
}()
i.InteractiveShell = &InteractiveShell{
StdoutCh: stdoutCh,
StderrCh: stderrCh,
Stdin: stdin,
Cmd: cmd,
}
err = cmd.Start()
if err != nil {
return err
}
return nil
}
func (i *Item) executeInteractiveShellCmd(variables []string) (string, error) {
_, err := io.WriteString(i.InteractiveShell.Stdin, i.SampleScript+"\n")
if err != nil {
return "", err
}
timeout := make(chan bool, 1)
go func() {
time.Sleep(time.Duration(i.RateMs / 2))
timeout <- true
}()
var resultText strings.Builder
var errorText strings.Builder
for {
select {
case stdout := <-i.InteractiveShell.StdoutCh:
if len(stdout) > 0 {
resultText.WriteString(stdout)
resultText.WriteString("\n")
}
case stderr := <-i.InteractiveShell.StderrCh:
if len(stderr) > 0 {
errorText.WriteString(stderr)
errorText.WriteString("\n")
}
case <-timeout:
if errorText.Len() > 0 {
return "", errors.New(errorText.String())
} else {
return i.transformInteractiveShellCmd(resultText.String())
}
}
}
}
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)
}
}