From 2596f6b0cd3c6ce5e88a0cf58919b1bfcf504fbf Mon Sep 17 00:00:00 2001 From: sqshq Date: Tue, 4 Jun 2019 23:24:16 -0400 Subject: [PATCH] added pty interactive shell setting, user now is able to choose between basic interactive shell and PTY interactive shell --- component/runchart/runchart.go | 4 +- config/component.go | 1 + config/default.go | 43 +++++++- data/int_pty.go | 129 ++++++++++++++++++++++++ data/int_shell.go | 110 +++++++++++++++++++++ data/item.go | 175 ++++++--------------------------- data/item_test.go | 25 ----- example-interactive-shell.yml | 3 + 8 files changed, 312 insertions(+), 178 deletions(-) create mode 100644 data/int_pty.go create mode 100644 data/int_shell.go delete mode 100644 data/item_test.go diff --git a/component/runchart/runchart.go b/component/runchart/runchart.go index 9174015..bc4a232 100644 --- a/component/runchart/runchart.go +++ b/component/runchart/runchart.go @@ -133,13 +133,13 @@ func (c *RunChart) Draw(buffer *ui.Buffer) { c.grid = c.newChartGrid() drawArea := image.Rect( - c.Inner.Min.X+c.grid.minTimeWidth+1, c.Inner.Min.Y, + c.Inner.Min.X+c.grid.minTimeWidth+2, c.Inner.Min.Y, c.Inner.Max.X, c.Inner.Max.Y-xAxisLabelsHeight-1, ) + c.renderAxes(buffer) c.renderLines(buffer, drawArea) c.renderLegend(buffer, drawArea) - c.renderAxes(buffer) component.RenderAlert(c.alert, c.Rectangle, buffer) c.mutex.Unlock() } diff --git a/config/component.go b/config/component.go index 244263a..b11fa27 100644 --- a/config/component.go +++ b/config/component.go @@ -104,6 +104,7 @@ type LegendConfig struct { type Item struct { Label *string `yaml:"label,omitempty"` Color *ui.Color `yaml:"color,omitempty"` + Pty *bool `yaml:"pty,omitempty"` InitScript *string `yaml:"init,omitempty"` SampleScript *string `yaml:"sample"` TransformScript *string `yaml:"transform,omitempty"` diff --git a/config/default.go b/config/default.go index be4fa8c..3fdf7fd 100644 --- a/config/default.go +++ b/config/default.go @@ -12,7 +12,7 @@ const ( func (c *Config) setDefaults() { c.setDefaultValues() - c.setDefaultColors() + c.setDefaultItemSettings() c.setDefaultArrangement() } @@ -173,17 +173,21 @@ func setDefaultTriggersValues(triggers []TriggerConfig) { } } -func (c *Config) setDefaultColors() { +func (c *Config) setDefaultItemSettings() { palette := console.GetPalette(*c.Theme) colorsCount := len(palette.ContentColors) + defaultPty := false for _, ch := range c.RunCharts { for j, item := range ch.Items { if item.Color == nil { item.Color = &palette.ContentColors[j%colorsCount] - ch.Items[j] = item } + if item.Pty == nil { + item.Pty = &defaultPty + } + ch.Items[j] = item } } @@ -191,20 +195,49 @@ func (c *Config) setDefaultColors() { for j, item := range b.Items { if item.Color == nil { item.Color = &palette.ContentColors[j%colorsCount] - b.Items[j] = item } + if item.Pty == nil { + item.Pty = &defaultPty + } + b.Items[j] = item } } for i, s := range c.SparkLines { s.Gradient = &palette.GradientColors[i%(len(palette.GradientColors))] + if s.Item.Pty == nil { + s.Item.Pty = &defaultPty + } c.SparkLines[i] = s } for i, g := range c.Gauges { + if g.Min.Pty == nil { + g.Min.Pty = &defaultPty + } + if g.Max.Pty == nil { + g.Max.Pty = &defaultPty + } + if g.Cur.Pty == nil { + g.Cur.Pty = &defaultPty + } if g.Color == nil { g.Color = &palette.ContentColors[i%colorsCount] - c.Gauges[i] = g } + c.Gauges[i] = g + } + + for i, a := range c.AsciiBoxes { + if a.Item.Pty == nil { + a.Item.Pty = &defaultPty + } + c.AsciiBoxes[i] = a + } + + for i, t := range c.TextBoxes { + if t.Item.Pty == nil { + t.Item.Pty = &defaultPty + } + c.TextBoxes[i] = t } } diff --git a/data/int_pty.go b/data/int_pty.go new file mode 100644 index 0000000..4c623f8 --- /dev/null +++ b/data/int_pty.go @@ -0,0 +1,129 @@ +package data + +import ( + "bufio" + "errors" + "fmt" + "github.com/kr/pty" + "github.com/lunixbochs/vtclean" + "io" + "os/exec" + "strings" + "time" +) + +const ( + startupTimeout = 100 * time.Millisecond + minAwaitTimeout = 100 * time.Millisecond + maxAwaitTimeout = 1 * time.Second +) + +/** + * Experimental + */ +type PtyInteractiveShell struct { + item *Item + variables []string + cmd *exec.Cmd + File io.WriteCloser + ch chan string + errCount int +} + +func (s *PtyInteractiveShell) init() error { + + cmd := exec.Command("sh", "-c", *s.item.initScript) + enrichEnvVariables(cmd, s.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() + } + }() + + s.cmd = cmd + s.File = file + s.ch = channel + + _, err = file.Read(make([]byte, 4096)) + if err != nil { + return err + } + + time.Sleep(startupTimeout) + + return nil +} + +func (s *PtyInteractiveShell) execute() (string, error) { + + _, err := io.WriteString(s.File, fmt.Sprintf(" %s\n", s.item.sampleScript)) + if err != nil { + s.errCount++ + if s.errCount > errorThreshold { + s.item.ptyShell = nil // restart session + } + return "", errors.New(fmt.Sprintf("Failed to execute command: %s", err)) + } + + softTimeout := make(chan bool, 1) + hardTimeout := make(chan bool, 1) + + go func() { + time.Sleep(s.getAwaitTimeout() / 2) + softTimeout <- true + time.Sleep(s.getAwaitTimeout() * 100) + hardTimeout <- true + }() + + var builder strings.Builder + softTimeoutElapsed := false + +await: + for { + select { + case out := <-s.ch: + cout := vtclean.Clean(out, false) + if len(cout) > 0 && !strings.Contains(cout, s.item.sampleScript) { + builder.WriteString(cout) + builder.WriteString("\n") + if softTimeoutElapsed { + break await + } + } + case <-softTimeout: + if builder.Len() > 0 { + break await + } else { + softTimeoutElapsed = true + } + case <-hardTimeout: + break await + } + } + + sample := strings.TrimSpace(builder.String()) + + return s.item.transform(sample) +} + +func (s *PtyInteractiveShell) getAwaitTimeout() time.Duration { + + timeout := time.Duration(s.item.rateMs) * time.Millisecond + + if timeout > maxAwaitTimeout { + return maxAwaitTimeout + } else if timeout < minAwaitTimeout { + return minAwaitTimeout + } + + return timeout +} diff --git a/data/int_shell.go b/data/int_shell.go new file mode 100644 index 0000000..1bfd8f5 --- /dev/null +++ b/data/int_shell.go @@ -0,0 +1,110 @@ +package data + +import ( + "bufio" + "errors" + "fmt" + "io" + "os/exec" + "strings" + "time" +) + +type BasicInteractiveShell struct { + item *Item + variables []string + stdoutCh chan string + stderrCh chan string + stdin io.WriteCloser + cmd *exec.Cmd + errCount int +} + +func (s *BasicInteractiveShell) init() error { + + cmd := exec.Command("sh", "-c", *s.item.initScript) + enrichEnvVariables(cmd, s.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() + } + }() + + s.stdoutCh = stdoutCh + s.stderrCh = stderrCh + s.stdin = stdin + s.cmd = cmd + + err = cmd.Start() + if err != nil { + return err + } + + return nil +} + +func (s *BasicInteractiveShell) execute() (string, error) { + + _, err := io.WriteString(s.stdin, s.item.sampleScript+"\n") + if err != nil { + s.errCount++ + if s.errCount > errorThreshold { + s.item.ptyShell = nil // restart session + } + return "", errors.New(fmt.Sprintf("Failed to execute command: %s", err)) + } + + timeout := make(chan bool, 1) + + go func() { + time.Sleep(time.Duration(s.item.rateMs / 2)) + timeout <- true + }() + + var resultText strings.Builder + var errorText strings.Builder + + for { + select { + case stdout := <-s.stdoutCh: + if len(stdout) > 0 { + resultText.WriteString(stdout) + resultText.WriteString("\n") + } + case stderr := <-s.stderrCh: + if len(stderr) > 0 { + errorText.WriteString(stderr) + errorText.WriteString("\n") + } + case <-timeout: + if errorText.Len() > 0 { + return "", errors.New(errorText.String()) + } else { + return resultText.String(), nil + } + } + } +} diff --git a/data/item.go b/data/item.go index f9315dd..aa04f8d 100644 --- a/data/item.go +++ b/data/item.go @@ -1,42 +1,24 @@ package data import ( - "bufio" - "errors" - "fmt" ui "github.com/gizak/termui/v3" - "github.com/kr/pty" - "github.com/lunixbochs/vtclean" "github.com/sqshq/sampler/config" - "io" "os" "os/exec" - "strings" - "time" ) -const ( - interactiveShellStartupTimeout = 100 * time.Millisecond - interactiveShellMinAwaitTimeout = 100 * time.Millisecond - interactiveShellMaxAwaitTimeout = 1 * time.Second - interactiveShellErrorThreshold = 10 -) +const errorThreshold = 10 type Item struct { - label string - sampleScript string - initScript *string - transformScript *string - color *ui.Color - rateMs int - errorsCount int - interactiveShell *InteractiveShell -} - -type InteractiveShell struct { - Channel chan string - File io.WriteCloser - Cmd *exec.Cmd + label string + sampleScript string + initScript *string + transformScript *string + color *ui.Color + rateMs int + pty bool + basicShell *BasicInteractiveShell + ptyShell *PtyInteractiveShell } func NewItems(cfgs []config.Item, rateMs int) []*Item { @@ -51,30 +33,32 @@ func NewItems(cfgs []config.Item, rateMs int) []*Item { 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 i.initScript != nil && i.interactiveShell == nil { + if i.initScript != nil && i.basicShell == nil && i.ptyShell == nil { err := i.initInteractiveShell(variables) if err != nil { - return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err)) + return "", err } } - if i.initScript != nil { - return i.executeInteractiveShellCmd(variables) + if i.basicShell != nil { + return i.basicShell.execute() + } else if i.ptyShell != nil { + return i.ptyShell.execute() } else { - return i.executeCmd(variables, i.sampleScript) + return i.execute(variables, i.sampleScript) } } -func (i *Item) executeCmd(variables []string, script string) (string, error) { +func (i *Item) execute(variables []string, script string) (string, error) { cmd := exec.Command("sh", "-c", script) enrichEnvVariables(cmd, variables) @@ -85,132 +69,31 @@ func (i *Item) executeCmd(variables []string, script string) (string, error) { return "", err } - result := vtclean.Clean(string(output), false) - - return result, nil + 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 +func (i *Item) initInteractiveShell(v []string) error { + if i.pty { + i.ptyShell = &PtyInteractiveShell{item: i, variables: v} + return i.ptyShell.init() + } else { + i.basicShell = &BasicInteractiveShell{item: i, variables: v} + return i.basicShell.init() } - - 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 { - i.errorsCount++ - if i.errorsCount > interactiveShellErrorThreshold { - i.interactiveShell = nil // restart session - i.errorsCount = 0 - } - return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err)) - } - - softTimeout := make(chan bool, 1) - hardTimeout := make(chan bool, 1) - - go func() { - time.Sleep(i.getAwaitTimeout() / 4) - softTimeout <- true - time.Sleep(i.getAwaitTimeout() * 100) - hardTimeout <- true - }() - - var builder strings.Builder - softTimeoutElapsed := false - -await: - for { - select { - case output := <-i.interactiveShell.Channel: - o := vtclean.Clean(output, false) - if len(o) > 0 && !strings.Contains(o, i.sampleScript) { - builder.WriteString(o) - builder.WriteString("\n") - if softTimeoutElapsed { - break await - } - } - case <-softTimeout: - if builder.Len() > 0 { - break await - } else { - softTimeoutElapsed = true - } - case <-hardTimeout: - break await - } - } - - sample := strings.TrimSpace(builder.String()) - - return i.transformInteractiveShellCmd(sample) -} - -func (i *Item) transformInteractiveShellCmd(sample string) (string, error) { +func (i *Item) transform(sample string) (string, error) { if i.transformScript != nil && len(sample) > 0 { - return i.executeCmd([]string{"sample=" + sample}, *i.transformScript) + return i.execute([]string{"sample=" + sample}, *i.transformScript) } return sample, nil } -func (i *Item) getAwaitTimeout() time.Duration { - - timeout := time.Duration(i.rateMs) * time.Millisecond - - if timeout > interactiveShellMaxAwaitTimeout { - return interactiveShellMaxAwaitTimeout - } else if timeout < interactiveShellMinAwaitTimeout { - return interactiveShellMinAwaitTimeout - } - - return timeout -} - 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 -} diff --git a/data/item_test.go b/data/item_test.go deleted file mode 100644 index 88cce01..0000000 --- a/data/item_test.go +++ /dev/null @@ -1,25 +0,0 @@ -package data - -import "testing" - -func Test_cleanupOutput(t *testing.T) { - type args struct { - output string - } - tests := []struct { - name string - args args - want string - }{ - {"should trim everything before carriage return", args{">>\rtext"}, "text"}, - {"should trim carriage return at the end", args{"text\r"}, "text"}, - {"should remove tabs and spaces", args{"\t\t\ntext "}, "text"}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := cleanupOutput(tt.args.output); got != tt.want { - t.Errorf("cleanupOutput() = %v, want %v", got, tt.want) - } - }) - } -} diff --git a/example-interactive-shell.yml b/example-interactive-shell.yml index cc7fe3a..2056fed 100644 --- a/example-interactive-shell.yml +++ b/example-interactive-shell.yml @@ -8,6 +8,7 @@ variables: textboxes: - title: Neo4j position: [[0, 0], [10, 40]] + pty: true init: $neo4jconnection sample: RETURN rand(); transform: echo "$sample" | tail -n 1 @@ -17,6 +18,7 @@ textboxes: sample: select random(); - title: MySQL position: [[19, 0], [10, 40]] + pty: true init: $mysqlconnection sample: select rand(); - title: MongoDB @@ -26,5 +28,6 @@ textboxes: sample: sleep(3000);Date.now(); - title: SSH position: [[39, 0], [41, 40]] + pty: true init: $sshconnection sample: top