forked from owenthereal/osh
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
181 lines (150 loc) · 3.57 KB
/
main.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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
package main
import (
"bufio"
"bytes"
"fmt"
"os"
"os/exec"
"regexp"
"strconv"
"strings"
)
type Builtin struct {
Run func(args []string)
}
var BUILTINS = map[string]*Builtin{
"cd": &Builtin{func(args []string) {
os.Chdir(args[0])
}},
"exit": &Builtin{func(args []string) {
var code int
if len(args) == 1 {
code, _ = strconv.Atoi(args[0])
}
os.Exit(code)
}},
"exec": &Builtin{func(args []string) {
cmd := exec.Command(args[0], args[1:]...)
spawnPrograms(cmd)
}},
"set": &Builtin{func(args []string) {
for _, arg := range args {
keyValuePair := strings.Split(arg, "=")
if len(keyValuePair) == 2 {
os.Setenv(keyValuePair[0], keyValuePair[1])
}
}
}},
}
func init() {
os.Setenv("PROMPT", "-> ")
}
func main() {
prompt()
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
commands := splitOnPipes(scanner.Text())
var cmds []*exec.Cmd
for _, command := range commands {
name, args := parseCommand(command)
if name == "" {
continue
}
if isBuiltin(name) {
callBuiltin(name, args)
} else {
cmd := exec.Command(name, args...)
cmds = append(cmds, cmd)
}
}
if len(cmds) > 0 {
spawnPrograms(cmds...)
}
prompt()
}
if err := scanner.Err(); err != nil {
fmt.Fprintln(os.Stderr, "Reading standard input:", err)
}
}
func prompt() {
fmt.Print(os.Getenv("PROMPT"))
}
func splitOnPipes(line string) (commands []string) {
pipeRegexp := regexp.MustCompile("([^\"'|]+)|[\"']([^\"']+)[\"']")
if pipeRegexp.MatchString(line) {
commands = pipeRegexp.FindAllString(line, -1)
} else {
commands = append(commands, line)
}
for i, command := range commands {
commands[i] = strings.TrimSpace(command)
}
return
}
func parseCommand(line string) (name string, args []string) {
regexpBySpace := regexp.MustCompile("\\s+")
cmd := regexpBySpace.Split(line, -1)
name = cmd[0]
// expand environment variables
// somehow os/exec.Command.Run() doesn't expand automatically
envVarRegexp := regexp.MustCompile("^\\$(.+)$")
for _, arg := range cmd[1:] {
if envVarRegexp.MatchString(arg) {
match := envVarRegexp.FindStringSubmatch(arg)
arg = os.Getenv(match[1])
}
args = append(args, arg)
}
return
}
func isBuiltin(name string) bool {
_, ok := BUILTINS[name]
return ok
}
func callBuiltin(name string, args []string) {
builtin, _ := BUILTINS[name]
builtin.Run(args)
}
func spawnPrograms(cmds ...*exec.Cmd) {
stdout, stderr, err := pipeline(cmds)
if err != nil {
fmt.Printf("%s\n", err)
}
if len(stdout) > 0 {
fmt.Print(string(stdout))
}
if len(stderr) > 0 {
fmt.Print(string(stderr))
}
}
func pipeline(cmds []*exec.Cmd) (pipeLineOutput, collectedStandardError []byte, pipeLineError error) {
// Collect the output from the command(s)
var output bytes.Buffer
var stderr bytes.Buffer
last := len(cmds) - 1
for i, cmd := range cmds[:last] {
var err error
// Connect each command's stdin to the previous command's stdout
if cmds[i+1].Stdin, err = cmd.StdoutPipe(); err != nil {
return nil, nil, err
}
// Connect each command's stderr to a buffer
cmd.Stderr = &stderr
}
// Connect the output and error for the last command
cmds[last].Stdout, cmds[last].Stderr = &output, &stderr
// Start each command
for _, cmd := range cmds {
if err := cmd.Start(); err != nil {
return output.Bytes(), stderr.Bytes(), err
}
}
// Wait for each command to complete
for _, cmd := range cmds {
if err := cmd.Wait(); err != nil {
return output.Bytes(), stderr.Bytes(), err
}
}
// Return the pipeline output and the collected standard error
return output.Bytes(), stderr.Bytes(), nil
}