forked from codeskyblue/go-sh
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pipe.go
130 lines (119 loc) · 2.33 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
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 = 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
}