/
command.go
131 lines (115 loc) · 3.43 KB
/
command.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package issh
import (
"context"
"errors"
"fmt"
"strconv"
"strings"
"time"
)
const (
// DefaultTimeOut archeve, Command is canneled. You can change by WithTimeoutOption
DefaultTimeOut = time.Second * 5
)
var (
// DefaultCallback is called after Command and just sleep in a second. You can change by WithCallbackOption
DefaultCallback = func(c *Command) (bool, error) {
time.Sleep(time.Millisecond * 500)
if c.Result.ReturnCode != 0 {
return false, fmt.Errorf("cmd [%v] %v", c.Input, ErrReturnCodeNotZero)
}
return true, nil
}
// ErrReturnCodeNotZero is error in command exit with non zero
ErrReturnCodeNotZero = errors.New("return code is not 0")
)
// Command has Input config and Output in remote host.
// Input is line of command execute in remote host.
// Callback is called after input command is finished. You can check whether Output is exepected in this function.
// NextCommand is called after Callback and called only Callback returns "true". NextCommand cannot has another NextCommand.
// ReturnCodeCheck is "true", Input is added ";echo $?" and check after Output is 0. Also you can manage retrun code in Callback.
// OutputLevel is logging level of command. Secret command should be set Silent
// Result is Command Output. You can use this in Callback, NextCommand, DefaultNextCommand functions.
type Command struct {
Input string
Callback func(c *Command) (bool, error)
NextCommand func(c *Command) *Command
ReturnCodeCheck bool
OutputLevel OutputLevel
Timeout time.Duration
Result *CommandResult
}
// CommandResult has command output and return code in remote host
type CommandResult struct {
Output []string
Lines int
ReturnCode int
}
// OutputLevel set logging level of command
type OutputLevel int
const (
// Silent logs nothing
Silent OutputLevel = iota
// Info logs only start and end of command
Info
// Output logs command output in remote host
Output
)
// NewCommand return Command with given options
func NewCommand(input string, options ...Option) *Command {
c := &Command{
Input: input,
OutputLevel: Info,
Timeout: DefaultTimeOut,
Callback: DefaultCallback,
ReturnCodeCheck: true,
Result: &CommandResult{},
}
for _, opt := range options {
opt.Apply(c)
}
if c.ReturnCodeCheck {
c.Input += ";echo $?"
}
return c
}
func (c *Command) wait(ctx context.Context, out <-chan string) error {
timeout, cancel := context.WithTimeout(ctx, c.Timeout)
defer cancel()
for {
select {
case v := <-out:
c.Result.Output = strings.Split(v, "\r\n")
c.Result.Lines = len(c.Result.Output)
if c.ReturnCodeCheck {
if c.Result.Lines-2 < 0 {
return fmt.Errorf("Couldn't check return code lines not enough %v", c.Result.Lines)
}
returnCode, err := strconv.Atoi(c.Result.Output[c.Result.Lines-2])
if err != nil {
return fmt.Errorf("Couldn't check retrun code %v", err)
}
c.Result.ReturnCode = returnCode
}
return nil
case <-timeout.Done():
if c.OutputLevel == Silent {
return errors.New("canceled by timeout or by parent")
}
return fmt.Errorf("[%v] is canceled by timeout or by parent", c.Input)
}
}
}
func (c *Command) output() ([]string, bool) {
if c.OutputLevel != Output {
return nil, false
}
var output []string
for i := 0; i < c.Result.Lines-1; i++ {
if c.Result.Output[i] == "0" {
break
}
output = append(output, c.Result.Output[i])
}
return output, true
}