sampler-fork/data/item.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
}