1package shell
2
3import (
4 "context"
5 "strings"
6 "testing"
7
8 "github.com/stretchr/testify/require"
9)
10
11func TestCommandBlocking(t *testing.T) {
12 tests := []struct {
13 name string
14 blockFuncs []BlockFunc
15 command string
16 shouldBlock bool
17 }{
18 {
19 name: "block simple command",
20 blockFuncs: []BlockFunc{
21 func(args []string) bool {
22 return len(args) > 0 && args[0] == "curl"
23 },
24 },
25 command: "curl https://example.com",
26 shouldBlock: true,
27 },
28 {
29 name: "allow non-blocked command",
30 blockFuncs: []BlockFunc{
31 func(args []string) bool {
32 return len(args) > 0 && args[0] == "curl"
33 },
34 },
35 command: "echo hello",
36 shouldBlock: false,
37 },
38 {
39 name: "block subcommand",
40 blockFuncs: []BlockFunc{
41 func(args []string) bool {
42 return len(args) >= 2 && args[0] == "brew" && args[1] == "install"
43 },
44 },
45 command: "brew install wget",
46 shouldBlock: true,
47 },
48 {
49 name: "allow different subcommand",
50 blockFuncs: []BlockFunc{
51 func(args []string) bool {
52 return len(args) >= 2 && args[0] == "brew" && args[1] == "install"
53 },
54 },
55 command: "brew list",
56 shouldBlock: false,
57 },
58 {
59 name: "block npm global install with -g",
60 blockFuncs: []BlockFunc{
61 ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
62 },
63 command: "npm install -g typescript",
64 shouldBlock: true,
65 },
66 {
67 name: "block npm global install with --global",
68 blockFuncs: []BlockFunc{
69 ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
70 },
71 command: "npm install --global typescript",
72 shouldBlock: true,
73 },
74 {
75 name: "allow npm local install",
76 blockFuncs: []BlockFunc{
77 ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
78 ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
79 },
80 command: "npm install typescript",
81 shouldBlock: false,
82 },
83 }
84
85 for _, tt := range tests {
86 t.Run(tt.name, func(t *testing.T) {
87 // Create a temporary directory for each test
88 tmpDir := t.TempDir()
89
90 shell := NewShell(&Options{
91 WorkingDir: tmpDir,
92 BlockFuncs: tt.blockFuncs,
93 })
94
95 _, _, err := shell.Exec(context.Background(), tt.command)
96
97 if tt.shouldBlock {
98 if err == nil {
99 t.Errorf("Expected command to be blocked, but it was allowed")
100 } else if !strings.Contains(err.Error(), "not allowed for security reasons") {
101 t.Errorf("Expected security error, got: %v", err)
102 }
103 } else {
104 // For non-blocked commands, we might get other errors (like command not found)
105 // but we shouldn't get the security error
106 if err != nil && strings.Contains(err.Error(), "not allowed for security reasons") {
107 t.Errorf("Command was unexpectedly blocked: %v", err)
108 }
109 }
110 })
111 }
112}
113
114func TestArgumentsBlocker(t *testing.T) {
115 tests := []struct {
116 name string
117 cmd string
118 args []string
119 flags []string
120 input []string
121 shouldBlock bool
122 }{
123 // Basic command blocking
124 {
125 name: "block exact command match",
126 cmd: "npm",
127 args: []string{"install"},
128 flags: nil,
129 input: []string{"npm", "install", "package"},
130 shouldBlock: true,
131 },
132 {
133 name: "allow different command",
134 cmd: "npm",
135 args: []string{"install"},
136 flags: nil,
137 input: []string{"yarn", "install", "package"},
138 shouldBlock: false,
139 },
140 {
141 name: "allow different subcommand",
142 cmd: "npm",
143 args: []string{"install"},
144 flags: nil,
145 input: []string{"npm", "list"},
146 shouldBlock: false,
147 },
148
149 // Flag-based blocking
150 {
151 name: "block with single flag",
152 cmd: "npm",
153 args: []string{"install"},
154 flags: []string{"-g"},
155 input: []string{"npm", "install", "-g", "typescript"},
156 shouldBlock: true,
157 },
158 {
159 name: "block with flag in different position",
160 cmd: "npm",
161 args: []string{"install"},
162 flags: []string{"-g"},
163 input: []string{"npm", "install", "typescript", "-g"},
164 shouldBlock: true,
165 },
166 {
167 name: "allow without required flag",
168 cmd: "npm",
169 args: []string{"install"},
170 flags: []string{"-g"},
171 input: []string{"npm", "install", "typescript"},
172 shouldBlock: false,
173 },
174 {
175 name: "block with multiple flags",
176 cmd: "pip",
177 args: []string{"install"},
178 flags: []string{"--user"},
179 input: []string{"pip", "install", "--user", "--upgrade", "package"},
180 shouldBlock: true,
181 },
182
183 // Complex argument patterns
184 {
185 name: "block multi-arg subcommand",
186 cmd: "yarn",
187 args: []string{"global", "add"},
188 flags: nil,
189 input: []string{"yarn", "global", "add", "typescript"},
190 shouldBlock: true,
191 },
192 {
193 name: "allow partial multi-arg match",
194 cmd: "yarn",
195 args: []string{"global", "add"},
196 flags: nil,
197 input: []string{"yarn", "global", "list"},
198 shouldBlock: false,
199 },
200
201 // Edge cases
202 {
203 name: "handle empty input",
204 cmd: "npm",
205 args: []string{"install"},
206 flags: nil,
207 input: []string{},
208 shouldBlock: false,
209 },
210 {
211 name: "handle command only",
212 cmd: "npm",
213 args: []string{"install"},
214 flags: nil,
215 input: []string{"npm"},
216 shouldBlock: false,
217 },
218 {
219 name: "block pacman with -S flag",
220 cmd: "pacman",
221 args: nil,
222 flags: []string{"-S"},
223 input: []string{"pacman", "-S", "package"},
224 shouldBlock: true,
225 },
226 {
227 name: "allow pacman without -S flag",
228 cmd: "pacman",
229 args: nil,
230 flags: []string{"-S"},
231 input: []string{"pacman", "-Q", "package"},
232 shouldBlock: false,
233 },
234
235 // `go test -exec`
236 {
237 name: "go test exec",
238 cmd: "go",
239 args: []string{"test"},
240 flags: []string{"-exec"},
241 input: []string{"go", "test", "-exec", "bash -c 'echo hello'"},
242 shouldBlock: true,
243 },
244 {
245 name: "go test exec",
246 cmd: "go",
247 args: []string{"test"},
248 flags: []string{"-exec"},
249 input: []string{"go", "test", `-exec="bash -c 'echo hello'"`},
250 shouldBlock: true,
251 },
252 }
253
254 for _, tt := range tests {
255 t.Run(tt.name, func(t *testing.T) {
256 blocker := ArgumentsBlocker(tt.cmd, tt.args, tt.flags)
257 result := blocker(tt.input)
258 require.Equal(t, tt.shouldBlock, result,
259 "Expected block=%v for input %v", tt.shouldBlock, tt.input)
260 })
261 }
262}
263
264func TestCommandsBlocker(t *testing.T) {
265 tests := []struct {
266 name string
267 banned []string
268 input []string
269 shouldBlock bool
270 }{
271 {
272 name: "block single banned command",
273 banned: []string{"curl"},
274 input: []string{"curl", "https://example.com"},
275 shouldBlock: true,
276 },
277 {
278 name: "allow non-banned command",
279 banned: []string{"curl", "wget"},
280 input: []string{"echo", "hello"},
281 shouldBlock: false,
282 },
283 {
284 name: "block from multiple banned",
285 banned: []string{"curl", "wget", "nc"},
286 input: []string{"wget", "https://example.com"},
287 shouldBlock: true,
288 },
289 {
290 name: "handle empty input",
291 banned: []string{"curl"},
292 input: []string{},
293 shouldBlock: false,
294 },
295 {
296 name: "case sensitive matching",
297 banned: []string{"curl"},
298 input: []string{"CURL", "https://example.com"},
299 shouldBlock: false,
300 },
301 }
302
303 for _, tt := range tests {
304 t.Run(tt.name, func(t *testing.T) {
305 blocker := CommandsBlocker(tt.banned)
306 result := blocker(tt.input)
307 require.Equal(t, tt.shouldBlock, result,
308 "Expected block=%v for input %v", tt.shouldBlock, tt.input)
309 })
310 }
311}
312
313func TestSplitArgsFlags(t *testing.T) {
314 tests := []struct {
315 name string
316 input []string
317 wantArgs []string
318 wantFlags []string
319 }{
320 {
321 name: "only args",
322 input: []string{"install", "package", "another"},
323 wantArgs: []string{"install", "package", "another"},
324 wantFlags: []string{},
325 },
326 {
327 name: "only flags",
328 input: []string{"-g", "--verbose", "-f"},
329 wantArgs: []string{},
330 wantFlags: []string{"-g", "--verbose", "-f"},
331 },
332 {
333 name: "mixed args and flags",
334 input: []string{"install", "-g", "package", "--verbose"},
335 wantArgs: []string{"install", "package"},
336 wantFlags: []string{"-g", "--verbose"},
337 },
338 {
339 name: "empty input",
340 input: []string{},
341 wantArgs: []string{},
342 wantFlags: []string{},
343 },
344 {
345 name: "single dash flag",
346 input: []string{"-S", "package"},
347 wantArgs: []string{"package"},
348 wantFlags: []string{"-S"},
349 },
350 {
351 name: "flag with equals sign",
352 input: []string{"-exec=bash", "package"},
353 wantArgs: []string{"package"},
354 wantFlags: []string{"-exec"},
355 },
356 {
357 name: "long flag with equals sign",
358 input: []string{"--config=/path/to/config", "run"},
359 wantArgs: []string{"run"},
360 wantFlags: []string{"--config"},
361 },
362 {
363 name: "flag with complex value",
364 input: []string{`-exec="bash -c 'echo hello'"`, "test"},
365 wantArgs: []string{"test"},
366 wantFlags: []string{"-exec"},
367 },
368 }
369
370 for _, tt := range tests {
371 t.Run(tt.name, func(t *testing.T) {
372 args, flags := splitArgsFlags(tt.input)
373 require.Equal(t, tt.wantArgs, args, "args mismatch")
374 require.Equal(t, tt.wantFlags, flags, "flags mismatch")
375 })
376 }
377}