/
pipe.go
148 lines (135 loc) · 2.63 KB
/
pipe.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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package sh
import (
"bytes"
"encoding/json"
"encoding/xml"
"errors"
"io"
"os"
"strings"
"syscall"
"time"
)
var ErrExecTimeout = errors.New("execute timeout")
// unmarshal shell output to decode json
func (s *Session) UnmarshalJSON(data interface{}) (err error) {
bufrw := bytes.NewBuffer(nil)
s.Stdout = bufrw
if err = s.Run(); err != nil {
return
}
return json.NewDecoder(bufrw).Decode(data)
}
// unmarshal command output into xml
func (s *Session) UnmarshalXML(data interface{}) (err error) {
bufrw := bytes.NewBuffer(nil)
s.Stdout = bufrw
if err = s.Run(); err != nil {
return
}
return xml.NewDecoder(bufrw).Decode(data)
}
// start command
func (s *Session) Start() (err error) {
s.started = true
var rd *io.PipeReader
var wr *io.PipeWriter
var length = len(s.cmds)
if s.ShowCMD {
var cmds = make([]string, 0, 4)
for _, cmd := range s.cmds {
cmds = append(cmds, strings.Join(cmd.Args, " "))
}
s.writePrompt(strings.Join(cmds, " | "))
}
for index, cmd := range s.cmds {
if index == 0 {
cmd.Stdin = s.Stdin
} else {
cmd.Stdin = rd
}
if index != length {
rd, wr = io.Pipe() // create pipe
cmd.Stdout = wr
cmd.Stderr = os.Stderr
}
if index == length-1 {
cmd.Stdout = s.Stdout
cmd.Stderr = s.Stderr
}
err = cmd.Start()
if err != nil {
return
}
}
return
}
// Should be call after Start()
// only catch the last command error
func (s *Session) Wait() (err error) {
for _, cmd := range s.cmds {
err = cmd.Wait()
wr, ok := cmd.Stdout.(*io.PipeWriter)
if ok {
wr.Close()
}
}
return err
}
func (s *Session) Kill(sig os.Signal) {
for _, cmd := range s.cmds {
if cmd.Process != nil {
cmd.Process.Signal(sig)
}
}
}
func (s *Session) WaitTimeout(timeout time.Duration) (err error) {
select {
case <-time.After(timeout):
s.Kill(syscall.SIGKILL)
return ErrExecTimeout
case err = <-Go(s.Wait):
return err
}
}
func Go(f func() error) chan error {
ch := make(chan error)
go func() {
ch <- f()
}()
return ch
}
func (s *Session) Run() (err error) {
if err = s.Start(); err != nil {
return
}
if s.timeout != time.Duration(0) {
return s.WaitTimeout(s.timeout)
}
return s.Wait()
}
func (s *Session) Output() (out []byte, err error) {
oldout := s.Stdout
defer func() {
s.Stdout = oldout
}()
stdout := bytes.NewBuffer(nil)
s.Stdout = stdout
err = s.Run()
out = stdout.Bytes()
return
}
func (s *Session) CombinedOutput() (out []byte, err error) {
oldout := s.Stdout
olderr := s.Stderr
defer func() {
s.Stdout = oldout
s.Stderr = olderr
}()
stdout := bytes.NewBuffer(nil)
s.Stdout = stdout
s.Stderr = stdout
err = s.Run()
out = stdout.Bytes()
return
}