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