git.iko.soy/crush
on main help sign in
fix(shell): refactor arguments blocker to check for flags in any position
commitcbdd2e522033a6286dce57addc9b6fd5a7d131d3 authorAndrey Nering <andreynering@users.noreply.github.com> dateMon 18 Aug 2025, 14:53 (1 year ago) parent532b473
3 files changed+288−53
compare with main
unifiedsplit
internal/llm/tools/bash.go+23−24 internal/shell/command_block_test.go+235−12 internal/shell/shell.go+30−17
internal/llm/tools/bash.go+23−24
@@ -275,30 +275,29 @@ Important:
275275func blockFuncs() []shell.BlockFunc {
276276 return []shell.BlockFunc{
277277 shell.CommandsBlocker(bannedCommands),
278− shell.ArgumentsBlocker([][]string{
279− // System package managers
280− {"apk", "add"},
281− {"apt", "install"},
282− {"apt-get", "install"},
283− {"dnf", "install"},
284− {"pacman", "-S"},
285− {"pkg", "install"},
286− {"yum", "install"},
287− {"zypper", "install"},
288−
289− // Language-specific package managers
290− {"brew", "install"},
291− {"cargo", "install"},
292− {"gem", "install"},
293− {"go", "install"},
294− {"npm", "install", "-g"},
295− {"npm", "install", "--global"},
296− {"pip", "install", "--user"},
297− {"pip3", "install", "--user"},
298− {"pnpm", "add", "-g"},
299− {"pnpm", "add", "--global"},
300− {"yarn", "global", "add"},
301− }),
278+
279+ // System package managers
280+ shell.ArgumentsBlocker("apk", []string{"add"}, nil),
281+ shell.ArgumentsBlocker("apt", []string{"install"}, nil),
282+ shell.ArgumentsBlocker("apt-get", []string{"install"}, nil),
283+ shell.ArgumentsBlocker("dnf", []string{"install"}, nil),
284+ shell.ArgumentsBlocker("pacman", nil, []string{"-S"}),
285+ shell.ArgumentsBlocker("pkg", []string{"install"}, nil),
286+ shell.ArgumentsBlocker("yum", []string{"install"}, nil),
287+ shell.ArgumentsBlocker("zypper", []string{"install"}, nil),
288+
289+ // Language-specific package managers
290+ shell.ArgumentsBlocker("brew", []string{"install"}, nil),
291+ shell.ArgumentsBlocker("cargo", []string{"install"}, nil),
292+ shell.ArgumentsBlocker("gem", []string{"install"}, nil),
293+ shell.ArgumentsBlocker("go", []string{"install"}, nil),
294+ shell.ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
295+ shell.ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
296+ shell.ArgumentsBlocker("pip", []string{"install"}, []string{"--user"}),
297+ shell.ArgumentsBlocker("pip3", []string{"install"}, []string{"--user"}),
298+ shell.ArgumentsBlocker("pnpm", []string{"add"}, []string{"--global"}),
299+ shell.ArgumentsBlocker("pnpm", []string{"add"}, []string{"-g"}),
300+ shell.ArgumentsBlocker("yarn", []string{"global", "add"}, nil),
302301 }
303302}
304303
internal/shell/command_block_test.go+235−12
@@ -4,6 +4,8 @@ import (
44 "context"
55 "strings"
66 "testing"
7+
8+ "github.com/stretchr/testify/require"
79)
810
911func TestCommandBlocking(t *testing.T) {
@@ -56,10 +58,7 @@ func TestCommandBlocking(t *testing.T) {
5658 {
5759 name: "block npm global install with -g",
5860 blockFuncs: []BlockFunc{
59− ArgumentsBlocker([][]string{
60− {"npm", "install", "-g"},
61− {"npm", "install", "--global"},
62− }),
61+ ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
6362 },
6463 command: "npm install -g typescript",
6564 shouldBlock: true,
@@ -67,10 +66,7 @@ func TestCommandBlocking(t *testing.T) {
6766 {
6867 name: "block npm global install with --global",
6968 blockFuncs: []BlockFunc{
70− ArgumentsBlocker([][]string{
71− {"npm", "install", "-g"},
72− {"npm", "install", "--global"},
73− }),
69+ ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
7470 },
7571 command: "npm install --global typescript",
7672 shouldBlock: true,
@@ -78,10 +74,8 @@ func TestCommandBlocking(t *testing.T) {
7874 {
7975 name: "allow npm local install",
8076 blockFuncs: []BlockFunc{
81− ArgumentsBlocker([][]string{
82− {"npm", "install", "-g"},
83− {"npm", "install", "--global"},
84− }),
77+ ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}),
78+ ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}),
8579 },
8680 command: "npm install typescript",
8781 shouldBlock: false,
@@ -116,3 +110,232 @@ func TestCommandBlocking(t *testing.T) {
116110 })
117111 }
118112}
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 (
1616 "errors"
1717 "fmt"
1818 "os"
19+ "slices"
1920 "strings"
2021 "sync"
2122
23+ "github.com/charmbracelet/crush/internal/slicesext"
2224 "mvdan.cc/sh/moreinterp/coreutils"
2325 "mvdan.cc/sh/v3/expand"
2426 "mvdan.cc/sh/v3/interp"
@@ -171,25 +173,36 @@ func CommandsBlocker(cmds []string) BlockFunc {
171173 }
172174}
173175
174−// ArgumentsBlocker creates a BlockFunc that blocks specific subcommands
175−func ArgumentsBlocker(blockedSubCommands [][]string) BlockFunc {
176− return func(args []string) bool {
177− for _, blocked := range blockedSubCommands {
178− if len(args) >= len(blocked) {
179− match := true
180− for i, part := range blocked {
181− if args[i] != part {
182− match = false
183− break
184− }
185− }
186− if match {
187− return true
188− }
189− }
176+// ArgumentsBlocker creates a BlockFunc that blocks specific subcommand
177+func ArgumentsBlocker(cmd string, args []string, flags []string) BlockFunc {
178+ return func(parts []string) bool {
179+ if len(parts) == 0 || parts[0] != cmd {
180+ return false
181+ }
182+
183+ argParts, flagParts := splitArgsFlags(parts[1:])
184+ if len(argParts) < len(args) || len(flagParts) < len(flags) {
185+ return false
186+ }
187+
188+ argsMatch := slices.Equal(argParts[:len(args)], args)
189+ flagsMatch := slicesext.IsSubset(flags, flagParts)
190+
191+ return argsMatch && flagsMatch
192+ }
193+}
194+
195+func splitArgsFlags(parts []string) (args []string, flags []string) {
196+ args = make([]string, 0, len(parts))
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)
190203 }
191− return false
192204 }
205+ return
193206}
194207
195208func (s *Shell) blockHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {