git.iko.soy/crush
on main help sign in
fix(shell): refactor arguments blocker to check for flags in any position
cbdd2e5Andrey Nering
internal/llm/tools/bash.go+23−24
compare with main
unifiedsplit
@@ -275,30 +275,29 @@ Important:
275func blockFuncs() []shell.BlockFunc {
275func blockFuncs() []shell.BlockFunc {
276 return []shell.BlockFunc{
276 return []shell.BlockFunc{
277 shell.CommandsBlocker(bannedCommands),
277 shell.CommandsBlocker(bannedCommands),
278− shell.ArgumentsBlocker([][]string{
278+
279− // System package managers
279+ // System package managers
280− {"apk", "add"},
280+ shell.ArgumentsBlocker("apk", []string{"add"}, nil),
281− {"apt", "install"},
281+ shell.ArgumentsBlocker("apt", []string{"install"}, nil),
282− {"apt-get", "install"},
282+ shell.ArgumentsBlocker("apt-get", []string{"install"}, nil),
283− {"dnf", "install"},
283+ shell.ArgumentsBlocker("dnf", []string{"install"}, nil),
284− {"pacman", "-S"},
284+ shell.ArgumentsBlocker("pacman", nil, []string{"-S"}),
285− {"pkg", "install"},
285+ shell.ArgumentsBlocker("pkg", []string{"install"}, nil),
286− {"yum", "install"},
286+ shell.ArgumentsBlocker("yum", []string{"install"}, nil),
287− {"zypper", "install"},
287+ shell.ArgumentsBlocker("zypper", []string{"install"}, nil),
288−
288+
289− // Language-specific package managers
289+ // Language-specific package managers
290− {"brew", "install"},
290+ shell.ArgumentsBlocker("brew", []string{"install"}, nil),
291− {"cargo", "install"},
291+ shell.ArgumentsBlocker("cargo", []string{"install"}, nil),
292− {"gem", "install"},
292+ shell.ArgumentsBlocker("gem", []string{"install"}, nil),
293− {"go", "install"},
293+ shell.ArgumentsBlocker("go", []string{"install"}, nil),
294− {"npm", "install", "-g"},
294+ shell.ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
295− {"npm", "install", "--global"},
295+ shell.ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
296− {"pip", "install", "--user"},
296+ shell.ArgumentsBlocker("pip", []string{"install"}, []string{"--user"}),
297− {"pip3", "install", "--user"},
297+ shell.ArgumentsBlocker("pip3", []string{"install"}, []string{"--user"}),
298− {"pnpm", "add", "-g"},
298+ shell.ArgumentsBlocker("pnpm", []string{"add"}, []string{"--global"}),
299− {"pnpm", "add", "--global"},
299+ shell.ArgumentsBlocker("pnpm", []string{"add"}, []string{"-g"}),
300− {"yarn", "global", "add"},
300+ shell.ArgumentsBlocker("yarn", []string{"global", "add"}, nil),
301− }),
302 }
301 }
303}
302}
304
303
internal/shell/command_block_test.go+235−12
@@ -4,6 +4,8 @@ import (
4 "context"
4 "context"
5 "strings"
5 "strings"
6 "testing"
6 "testing"
7+
8+ "github.com/stretchr/testify/require"
7)
9)
8
10
9func TestCommandBlocking(t *testing.T) {
11func TestCommandBlocking(t *testing.T) {
@@ -56,10 +58,7 @@ func TestCommandBlocking(t *testing.T) {
56 {
58 {
57 name: "block npm global install with -g",
59 name: "block npm global install with -g",
58 blockFuncs: []BlockFunc{
60 blockFuncs: []BlockFunc{
59− ArgumentsBlocker([][]string{
61+ ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
60− {"npm", "install", "-g"},
61− {"npm", "install", "--global"},
62− }),
63 },
62 },
64 command: "npm install -g typescript",
63 command: "npm install -g typescript",
65 shouldBlock: true,
64 shouldBlock: true,
@@ -67,10 +66,7 @@ func TestCommandBlocking(t *testing.T) {
67 {
66 {
68 name: "block npm global install with --global",
67 name: "block npm global install with --global",
69 blockFuncs: []BlockFunc{
68 blockFuncs: []BlockFunc{
70− ArgumentsBlocker([][]string{
69+ ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
71− {"npm", "install", "-g"},
72− {"npm", "install", "--global"},
73− }),
74 },
70 },
75 command: "npm install --global typescript",
71 command: "npm install --global typescript",
76 shouldBlock: true,
72 shouldBlock: true,
@@ -78,10 +74,8 @@ func TestCommandBlocking(t *testing.T) {
78 {
74 {
79 name: "allow npm local install",
75 name: "allow npm local install",
80 blockFuncs: []BlockFunc{
76 blockFuncs: []BlockFunc{
81− ArgumentsBlocker([][]string{
77+ ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
82− {"npm", "install", "-g"},
83− {"npm", "install", "--global"},
78+ ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
84− }),
85 },
79 },
86 command: "npm install typescript",
80 command: "npm install typescript",
87 shouldBlock: false,
81 shouldBlock: false,
@@ -116,3 +110,232 @@ func TestCommandBlocking(t *testing.T) {
116 })
110 })
117 }
111 }
118}
112}
113+
114+func 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+
236+ for _, tt := range tests {
237+ t.Run(tt.name, func(t *testing.T) {
238+ blocker := ArgumentsBlocker(tt.cmd, tt.args, tt.flags)
239+ result := blocker(tt.input)
240+ require.Equal(t, tt.shouldBlock, result,
241+ "Expected block=%v for input %v", tt.shouldBlock, tt.input)
242+ })
243+ }
244+}
245+
246+func TestCommandsBlocker(t *testing.T) {
247+ tests := []struct {
248+ name string
249+ banned []string
250+ input []string
251+ shouldBlock bool
252+ }{
253+ {
254+ name: "block single banned command",
255+ banned: []string{"curl"},
256+ input: []string{"curl", "https://example.com"},
257+ shouldBlock: true,
258+ },
259+ {
260+ name: "allow non-banned command",
261+ banned: []string{"curl", "wget"},
262+ input: []string{"echo", "hello"},
263+ shouldBlock: false,
264+ },
265+ {
266+ name: "block from multiple banned",
267+ banned: []string{"curl", "wget", "nc"},
268+ input: []string{"wget", "https://example.com"},
269+ shouldBlock: true,
270+ },
271+ {
272+ name: "handle empty input",
273+ banned: []string{"curl"},
274+ input: []string{},
275+ shouldBlock: false,
276+ },
277+ {
278+ name: "case sensitive matching",
279+ banned: []string{"curl"},
280+ input: []string{"CURL", "https://example.com"},
281+ shouldBlock: false,
282+ },
283+ }
284+
285+ for _, tt := range tests {
286+ t.Run(tt.name, func(t *testing.T) {
287+ blocker := CommandsBlocker(tt.banned)
288+ result := blocker(tt.input)
289+ require.Equal(t, tt.shouldBlock, result,
290+ "Expected block=%v for input %v", tt.shouldBlock, tt.input)
291+ })
292+ }
293+}
294+
295+func TestSplitArgsFlags(t *testing.T) {
296+ tests := []struct {
297+ name string
298+ input []string
299+ wantArgs []string
300+ wantFlags []string
301+ }{
302+ {
303+ name: "only args",
304+ input: []string{"install", "package", "another"},
305+ wantArgs: []string{"install", "package", "another"},
306+ wantFlags: []string{},
307+ },
308+ {
309+ name: "only flags",
310+ input: []string{"-g", "--verbose", "-f"},
311+ wantArgs: []string{},
312+ wantFlags: []string{"-g", "--verbose", "-f"},
313+ },
314+ {
315+ name: "mixed args and flags",
316+ input: []string{"install", "-g", "package", "--verbose"},
317+ wantArgs: []string{"install", "package"},
318+ wantFlags: []string{"-g", "--verbose"},
319+ },
320+ {
321+ name: "empty input",
322+ input: []string{},
323+ wantArgs: []string{},
324+ wantFlags: []string{},
325+ },
326+ {
327+ name: "single dash flag",
328+ input: []string{"-S", "package"},
329+ wantArgs: []string{"package"},
330+ wantFlags: []string{"-S"},
331+ },
332+ }
333+
334+ for _, tt := range tests {
335+ t.Run(tt.name, func(t *testing.T) {
336+ args, flags := splitArgsFlags(tt.input)
337+ require.Equal(t, tt.wantArgs, args, "args mismatch")
338+ require.Equal(t, tt.wantFlags, flags, "flags mismatch")
339+ })
340+ }
341+}
internal/shell/shell.go+30−17
@@ -16,9 +16,11 @@ import (
16 "errors"
16 "errors"
17 "fmt"
17 "fmt"
18 "os"
18 "os"
19+ "slices"
19 "strings"
20 "strings"
20 "sync"
21 "sync"
21
22
23+ "github.com/charmbracelet/crush/internal/slicesext"
22 "mvdan.cc/sh/moreinterp/coreutils"
24 "mvdan.cc/sh/moreinterp/coreutils"
23 "mvdan.cc/sh/v3/expand"
25 "mvdan.cc/sh/v3/expand"
24 "mvdan.cc/sh/v3/interp"
26 "mvdan.cc/sh/v3/interp"
@@ -171,25 +173,36 @@ func CommandsBlocker(cmds []string) BlockFunc {
171 }
173 }
172}
174}
173
175
174−// ArgumentsBlocker creates a BlockFunc that blocks specific subcommands
176+// ArgumentsBlocker creates a BlockFunc that blocks specific subcommand
175−func ArgumentsBlocker(blockedSubCommands [][]string) BlockFunc {
177+func ArgumentsBlocker(cmd string, args []string, flags []string) BlockFunc {
176− return func(args []string) bool {
178+ return func(parts []string) bool {
177− for _, blocked := range blockedSubCommands {
179+ if len(parts) == 0 || parts[0] != cmd {
178− if len(args) >= len(blocked) {
180+ return false
179− match := true
181+ }
180− for i, part := range blocked {
182+
181− if args[i] != part {
183+ argParts, flagParts := splitArgsFlags(parts[1:])
184+ if len(argParts) < len(args) || len(flagParts) < len(flags) {
182− match = false
185+ return false
183− break
186+ }
187+
188+ argsMatch := slices.Equal(argParts[:len(args)], args)
189+ flagsMatch := slicesext.IsSubset(flags, flagParts)
190+
191+ return argsMatch && flagsMatch
184− }
192+ }
185− }
193+}
186− if match {
194+
187− return true
195+func splitArgsFlags(parts []string) (args []string, flags []string) {
188− }
196+ args = make([]string, 0, len(parts))
189− }
197+ flags = make([]string, 0, len(parts))
198+ for _, part := range parts {
199+ if strings.HasPrefix(part, "-") {
200+ flags = append(flags, part)
201+ } else {
202+ args = append(args, part)
190 }
203 }
191− return false
192 }
204 }
205+ return
193}
206}
194
207
195func (s *Shell) blockHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
208func (s *Shell) blockHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {