1package shell
2
3import (
4 "context"
5 "os"
6 "strings"
7 "testing"
8)
9
10func TestCommandBlocking(t *testing.T) {
11 tests := []struct {
12 name string
13 blockFuncs []CommandBlockFunc
14 command string
15 shouldBlock bool
16 }{
17 {
18 name: "block simple command",
19 blockFuncs: []CommandBlockFunc{
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: []CommandBlockFunc{
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: []CommandBlockFunc{
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: []CommandBlockFunc{
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: []CommandBlockFunc{
60 CreateSubCommandBlocker([][]string{
61 {"npm", "install", "-g"},
62 {"npm", "install", "--global"},
63 }),
64 },
65 command: "npm install -g typescript",
66 shouldBlock: true,
67 },
68 {
69 name: "block npm global install with --global",
70 blockFuncs: []CommandBlockFunc{
71 CreateSubCommandBlocker([][]string{
72 {"npm", "install", "-g"},
73 {"npm", "install", "--global"},
74 }),
75 },
76 command: "npm install --global typescript",
77 shouldBlock: true,
78 },
79 {
80 name: "allow npm local install",
81 blockFuncs: []CommandBlockFunc{
82 CreateSubCommandBlocker([][]string{
83 {"npm", "install", "-g"},
84 {"npm", "install", "--global"},
85 }),
86 },
87 command: "npm install typescript",
88 shouldBlock: false,
89 },
90 }
91
92 for _, tt := range tests {
93 t.Run(tt.name, func(t *testing.T) {
94 // Create a temporary directory for each test
95 tmpDir, err := os.MkdirTemp("", "shell-test-*")
96 if err != nil {
97 t.Fatalf("Failed to create temp dir: %v", err)
98 }
99 defer os.RemoveAll(tmpDir)
100
101 shell := NewShell(&Options{
102 WorkingDir: tmpDir,
103 BlockFuncs: tt.blockFuncs,
104 })
105
106 _, _, err = shell.Exec(context.Background(), tt.command)
107
108 if tt.shouldBlock {
109 if err == nil {
110 t.Errorf("Expected command to be blocked, but it was allowed")
111 } else if !strings.Contains(err.Error(), "not allowed for security reasons") {
112 t.Errorf("Expected security error, got: %v", err)
113 }
114 } else {
115 // For non-blocked commands, we might get other errors (like command not found)
116 // but we shouldn't get the security error
117 if err != nil && strings.Contains(err.Error(), "not allowed for security reasons") {
118 t.Errorf("Command was unexpectedly blocked: %v", err)
119 }
120 }
121 })
122 }
123}