sampler-fork/data/item.go

118 lines
2.4 KiB
Go
Raw Normal View History

2019-01-31 23:40:05 +00:00
package data
import (
2019-03-14 03:01:44 +00:00
ui "github.com/gizak/termui/v3"
"github.com/lunixbochs/vtclean"
"github.com/sqshq/sampler/config"
2019-03-21 02:23:08 +00:00
"os"
2019-01-31 23:40:05 +00:00
"os/exec"
2019-06-21 04:32:02 +00:00
"time"
2019-01-31 23:40:05 +00:00
)
const errorThreshold = 10
2019-01-31 23:40:05 +00:00
type Item struct {
label string
2019-06-23 04:08:24 +00:00
initScripts []string
sampleScript string
transformScript *string
color *ui.Color
rateMs int
pty bool
basicShell InteractiveShell
ptyShell InteractiveShell
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,
2019-06-23 04:08:24 +00:00
initScripts: getInitScripts(i),
transformScript: i.TransformScript,
color: i.Color,
rateMs: rateMs,
pty: *i.Pty,
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-06-23 04:08:24 +00:00
if len(i.initScripts) > 0 && i.basicShell == nil && i.ptyShell == nil {
2019-04-07 15:09:24 +00:00
err := i.initInteractiveShell(variables)
if err != nil {
return "", err
2019-04-07 15:09:24 +00:00
}
}
2019-03-21 02:23:08 +00:00
if i.basicShell != nil {
return i.basicShell.execute()
} else if i.ptyShell != nil {
return i.ptyShell.execute()
2019-04-07 15:09:24 +00:00
} else {
return i.execute(variables, i.sampleScript)
2019-03-21 02:23:08 +00:00
}
2019-04-07 15:09:24 +00:00
}
func (i *Item) execute(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
}
o := string(output)
return vtclean.Clean(o, false), nil
2019-04-07 15:09:24 +00:00
}
func (i *Item) initInteractiveShell(v []string) error {
2019-07-27 04:15:35 +00:00
2019-06-21 04:32:02 +00:00
timeout := time.Duration(i.rateMs) * time.Millisecond * 3 / 4
2019-07-27 04:15:35 +00:00
if i.pty {
2019-06-21 04:32:02 +00:00
i.ptyShell = &PtyInteractiveShell{item: i, variables: v, timeout: timeout}
return i.ptyShell.init()
2019-04-07 15:09:24 +00:00
}
2019-07-27 04:15:35 +00:00
i.basicShell = &BasicInteractiveShell{item: i, variables: v, timeout: timeout}
return i.basicShell.init()
2019-04-07 15:09:24 +00:00
}
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
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)
}
}
2019-06-23 04:08:24 +00:00
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{}
}
}