package data import ( ui "github.com/gizak/termui/v3" "github.com/lunixbochs/vtclean" "github.com/sqshq/sampler/config" "os" "os/exec" "time" ) const errorThreshold = 10 type Item struct { label string initScripts []string sampleScript string transformScript *string color *ui.Color rateMs int pty bool basicShell InteractiveShell ptyShell InteractiveShell } func NewItems(cfgs []config.Item, rateMs int) []*Item { items := make([]*Item, 0) for _, i := range cfgs { item := &Item{ label: *i.Label, sampleScript: *i.SampleScript, initScripts: getInitScripts(i), transformScript: i.TransformScript, color: i.Color, rateMs: rateMs, pty: *i.Pty, } items = append(items, item) } return items } func (i *Item) nextValue(variables []string) (string, error) { if len(i.initScripts) > 0 && i.basicShell == nil && i.ptyShell == nil { err := i.initInteractiveShell(variables) if err != nil { return "", err } } if i.basicShell != nil { return i.basicShell.execute() } else if i.ptyShell != nil { return i.ptyShell.execute() } else { return i.execute(variables, i.sampleScript) } } func (i *Item) execute(variables []string, script string) (string, error) { cmd := exec.Command("sh", "-c", script) enrichEnvVariables(cmd, variables) output, err := cmd.Output() if err != nil { return "", err } o := string(output) return vtclean.Clean(o, false), nil } func (i *Item) initInteractiveShell(v []string) error { timeout := time.Duration(i.rateMs) * time.Millisecond * 3 / 4 if i.pty { i.ptyShell = &PtyInteractiveShell{item: i, variables: v, timeout: timeout} return i.ptyShell.init() } i.basicShell = &BasicInteractiveShell{item: i, variables: v, timeout: timeout} return i.basicShell.init() } func (i *Item) transform(sample string) (string, error) { if i.transformScript != nil && len(sample) > 0 { return i.execute([]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 getInitScripts(item config.Item) []string { if item.MultiStepInitScript != nil { return *item.MultiStepInitScript } else if item.InitScript != nil { return []string{*item.InitScript} } else { return []string{} } }