forked from k3s-io/go-powershell
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathshell.go
137 lines (103 loc) · 2.81 KB
/
shell.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
// Copyright (c) 2017 Gorillalabs. All rights reserved.
package powershell
import (
"fmt"
"io"
"regexp"
"sync"
"github.com/acarl005/stripansi"
"github.com/direktiv/go-powershell/backend"
"github.com/direktiv/go-powershell/utils"
"github.com/pkg/errors"
)
const newline = "\r\n"
type Shell interface {
Execute(cmd string) (string, string, error)
Exit()
}
type shell struct {
handle backend.Waiter
stdin io.Writer
stdout io.Reader
stderr io.Reader
be backend.Starter
}
func New(be backend.Starter) (Shell, error) {
handle, stdin, stdout, stderr, err := be.StartProcess("/bin/pwsh", "-NoExit", "-Command", "-")
if err != nil {
return nil, err
}
return &shell{handle, stdin, stdout, stderr, be}, nil
}
func (s *shell) Execute(cmd string) (string, string, error) {
if s.handle == nil {
return "", "", errors.Wrap(errors.New(cmd), "Cannot execute commands on closed shells.")
}
outBoundary := createBoundary()
errBoundary := createBoundary()
// wrap the command in special markers so we know when to stop reading from the pipes
full := fmt.Sprintf("%s; echo '%s'; [Console]::Error.WriteLine('%s')%s", cmd, outBoundary, errBoundary, newline)
_, err := s.stdin.Write([]byte(full))
if err != nil {
return "", "", errors.Wrap(errors.Wrap(err, cmd), "Could not send PowerShell command")
}
// read stdout and stderr
sout := ""
serr := ""
waiter := &sync.WaitGroup{}
waiter.Add(2)
var wr io.Writer
bl, ok := s.be.(*backend.Local)
if ok {
wr = bl.Writer
}
go streamReader(s.stdout, outBoundary, &sout, waiter, wr)
go streamReader(s.stderr, errBoundary, &serr, waiter, wr)
waiter.Wait()
if len(serr) > 0 {
return sout, serr, errors.Wrap(errors.New(cmd), serr)
}
return sout, serr, nil
}
func (s *shell) Exit() {
s.stdin.Write([]byte("exit" + newline))
// if it's possible to close stdin, do so (some backends, like the local one,
// do support it)
closer, ok := s.stdin.(io.Closer)
if ok {
closer.Close()
}
s.handle.Wait()
s.handle = nil
s.stdin = nil
s.stdout = nil
s.stderr = nil
}
func streamReader(stream io.Reader, boundary string, buffer *string, signal *sync.WaitGroup, wr io.Writer) error {
// read all output until we have found our boundary token
output := ""
bufsize := 64
marker := regexp.MustCompile("(?s)(.*)" + regexp.QuoteMeta(boundary))
marker2 := regexp.MustCompile("\\$gorilla[a-z0-9]*")
for {
buf := make([]byte, bufsize)
read, err := stream.Read(buf)
if err != nil {
return err
}
output = output + string(buf[:read])
if marker.MatchString(output) {
break
}
if wr != nil {
sb := marker2.ReplaceAll(buf[:read], []byte(""))
wr.Write([]byte(stripansi.Strip(string(sb))))
}
}
*buffer = marker.FindStringSubmatch(output)[1]
signal.Done()
return nil
}
func createBoundary() string {
return "$gorilla" + utils.CreateRandomString(12) + "$"
}