From 2fdbcac0290e1aa977de0139c335210ce35b0faf Mon Sep 17 00:00:00 2001 From: Kujtim Hoxha Date: Fri, 11 Jul 2025 17:55:01 +0200 Subject: [PATCH] chore: better naming --- internal/llm/tools/bash.go | 101 ++++++++++++++++++++++++--- internal/shell/command_block_test.go | 24 +++---- internal/shell/shell.go | 25 +++---- 3 files changed, 118 insertions(+), 32 deletions(-) diff --git a/internal/llm/tools/bash.go b/internal/llm/tools/bash.go index feac86e01cc7769d886f90a926ad20b87741702d..6d7a9a32b3829da02021be80e6e41e28888efd83 100644 --- a/internal/llm/tools/bash.go +++ b/internal/llm/tools/bash.go @@ -42,9 +42,74 @@ const ( ) var bannedCommands = []string{ - "alias", "curl", "curlie", "wget", "axel", "aria2c", - "nc", "telnet", "lynx", "w3m", "links", "httpie", "xh", - "http-prompt", "chrome", "firefox", "safari", "sudo", + // Network/Download tools + "alias", + "aria2c", + "axel", + "chrome", + "curl", + "curlie", + "firefox", + "http-prompt", + "httpie", + "links", + "lynx", + "nc", + "safari", + "telnet", + "w3m", + "wget", + "xh", + + // System administration + "doas", + "su", + "sudo", + + // Package managers + "apk", + "apt", + "apt-cache", + "apt-get", + "dnf", + "dpkg", + "emerge", + "home-manager", + "makepkg", + "opkg", + "pacman", + "paru", + "pkg", + "pkg_add", + "pkg_delete", + "portage", + "rpm", + "yay", + "yum", + "zypper", + + // System modification + "at", + "batch", + "chkconfig", + "crontab", + "fdisk", + "mkfs", + "mount", + "parted", + "service", + "systemctl", + "umount", + + // Network configuration + "firewall-cmd", + "ifconfig", + "ip", + "iptables", + "netstat", + "pfctl", + "route", + "ufw", } // getSafeReadOnlyCommands returns platform-appropriate safe commands @@ -245,13 +310,33 @@ Important: - Never update git config`, bannedCommandsStr, MaxOutputLength) } -func createCommandBlockFuncs() []shell.CommandBlockFunc { - return []shell.CommandBlockFunc{ - shell.CreateSimpleCommandBlocker(bannedCommands), - shell.CreateSubCommandBlocker([][]string{ +func blockFuncs() []shell.BlockFunc { + return []shell.BlockFunc{ + shell.CommandsBlocker(bannedCommands), + shell.ArgumentsBlocker([][]string{ + // System package managers + {"apk", "add"}, + {"apt", "install"}, + {"apt-get", "install"}, + {"dnf", "install"}, + {"emerge"}, + {"pacman", "-S"}, + {"pkg", "install"}, + {"yum", "install"}, + {"zypper", "install"}, + + // Language-specific package managers {"brew", "install"}, + {"cargo", "install"}, + {"gem", "install"}, + {"go", "install"}, {"npm", "install", "-g"}, {"npm", "install", "--global"}, + {"pip", "install", "--user"}, + {"pip3", "install", "--user"}, + {"pnpm", "add", "-g"}, + {"pnpm", "add", "--global"}, + {"yarn", "global", "add"}, }), } } @@ -259,7 +344,7 @@ func createCommandBlockFuncs() []shell.CommandBlockFunc { func NewBashTool(permission permission.Service, workingDir string) BaseTool { // Set up command blocking on the persistent shell persistentShell := shell.GetPersistentShell(workingDir) - persistentShell.SetBlockFuncs(createCommandBlockFuncs()) + persistentShell.SetBlockFuncs(blockFuncs()) return &bashTool{ permissions: permission, diff --git a/internal/shell/command_block_test.go b/internal/shell/command_block_test.go index 85971748f882cf79fba3ea86d2682ce6ce4f252d..fd7c46bcd98e54f44abbe982e834f3cbb04cbfa4 100644 --- a/internal/shell/command_block_test.go +++ b/internal/shell/command_block_test.go @@ -10,13 +10,13 @@ import ( func TestCommandBlocking(t *testing.T) { tests := []struct { name string - blockFuncs []CommandBlockFunc + blockFuncs []BlockFunc command string shouldBlock bool }{ { name: "block simple command", - blockFuncs: []CommandBlockFunc{ + blockFuncs: []BlockFunc{ func(args []string) bool { return len(args) > 0 && args[0] == "curl" }, @@ -26,7 +26,7 @@ func TestCommandBlocking(t *testing.T) { }, { name: "allow non-blocked command", - blockFuncs: []CommandBlockFunc{ + blockFuncs: []BlockFunc{ func(args []string) bool { return len(args) > 0 && args[0] == "curl" }, @@ -36,7 +36,7 @@ func TestCommandBlocking(t *testing.T) { }, { name: "block subcommand", - blockFuncs: []CommandBlockFunc{ + blockFuncs: []BlockFunc{ func(args []string) bool { return len(args) >= 2 && args[0] == "brew" && args[1] == "install" }, @@ -46,7 +46,7 @@ func TestCommandBlocking(t *testing.T) { }, { name: "allow different subcommand", - blockFuncs: []CommandBlockFunc{ + blockFuncs: []BlockFunc{ func(args []string) bool { return len(args) >= 2 && args[0] == "brew" && args[1] == "install" }, @@ -56,8 +56,8 @@ func TestCommandBlocking(t *testing.T) { }, { name: "block npm global install with -g", - blockFuncs: []CommandBlockFunc{ - CreateSubCommandBlocker([][]string{ + blockFuncs: []BlockFunc{ + ArgumentsBlocker([][]string{ {"npm", "install", "-g"}, {"npm", "install", "--global"}, }), @@ -67,8 +67,8 @@ func TestCommandBlocking(t *testing.T) { }, { name: "block npm global install with --global", - blockFuncs: []CommandBlockFunc{ - CreateSubCommandBlocker([][]string{ + blockFuncs: []BlockFunc{ + ArgumentsBlocker([][]string{ {"npm", "install", "-g"}, {"npm", "install", "--global"}, }), @@ -78,8 +78,8 @@ func TestCommandBlocking(t *testing.T) { }, { name: "allow npm local install", - blockFuncs: []CommandBlockFunc{ - CreateSubCommandBlocker([][]string{ + blockFuncs: []BlockFunc{ + ArgumentsBlocker([][]string{ {"npm", "install", "-g"}, {"npm", "install", "--global"}, }), @@ -120,4 +120,4 @@ func TestCommandBlocking(t *testing.T) { } }) } -} \ No newline at end of file +} diff --git a/internal/shell/shell.go b/internal/shell/shell.go index 815be0907a3fd05f996a24f84f751ce5d776b833..097af74d0172264efbe899711ff57137abc6ee30 100644 --- a/internal/shell/shell.go +++ b/internal/shell/shell.go @@ -44,8 +44,8 @@ type noopLogger struct{} func (noopLogger) InfoPersist(msg string, keysAndValues ...interface{}) {} -// CommandBlockFunc is a function that determines if a command should be blocked -type CommandBlockFunc func(args []string) bool +// BlockFunc is a function that determines if a command should be blocked +type BlockFunc func(args []string) bool // Shell provides cross-platform shell execution with optional state persistence type Shell struct { @@ -53,7 +53,7 @@ type Shell struct { cwd string mu sync.Mutex logger Logger - blockFuncs []CommandBlockFunc + blockFuncs []BlockFunc } // Options for creating a new shell @@ -61,7 +61,7 @@ type Options struct { WorkingDir string Env []string Logger Logger - BlockFuncs []CommandBlockFunc + BlockFuncs []BlockFunc } // NewShell creates a new shell instance with the given options @@ -159,7 +159,7 @@ func (s *Shell) SetEnv(key, value string) { } // SetBlockFuncs sets the command block functions for the shell -func (s *Shell) SetBlockFuncs(blockFuncs []CommandBlockFunc) { +func (s *Shell) SetBlockFuncs(blockFuncs []BlockFunc) { s.mu.Lock() defer s.mu.Unlock() s.blockFuncs = blockFuncs @@ -216,13 +216,13 @@ func (s *Shell) determineShellType(command string) ShellType { return ShellTypePOSIX } -// CreateSimpleCommandBlocker creates a CommandBlockFunc that blocks exact command matches -func CreateSimpleCommandBlocker(bannedCommands []string) CommandBlockFunc { +// CommandsBlocker creates a BlockFunc that blocks exact command matches +func CommandsBlocker(bannedCommands []string) BlockFunc { bannedSet := make(map[string]bool) for _, cmd := range bannedCommands { bannedSet[cmd] = true } - + return func(args []string) bool { if len(args) == 0 { return false @@ -231,8 +231,8 @@ func CreateSimpleCommandBlocker(bannedCommands []string) CommandBlockFunc { } } -// CreateSubCommandBlocker creates a CommandBlockFunc that blocks specific subcommands -func CreateSubCommandBlocker(blockedSubCommands [][]string) CommandBlockFunc { +// ArgumentsBlocker creates a BlockFunc that blocks specific subcommands +func ArgumentsBlocker(blockedSubCommands [][]string) BlockFunc { return func(args []string) bool { for _, blocked := range blockedSubCommands { if len(args) >= len(blocked) { @@ -251,7 +251,8 @@ func CreateSubCommandBlocker(blockedSubCommands [][]string) CommandBlockFunc { return false } } -func (s *Shell) createCommandBlockHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { + +func (s *Shell) blockHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { return func(ctx context.Context, args []string) error { if len(args) == 0 { @@ -357,7 +358,7 @@ func (s *Shell) execPOSIX(ctx context.Context, command string) (string, string, interp.Interactive(false), interp.Env(expand.ListEnviron(s.env...)), interp.Dir(s.cwd), - interp.ExecHandlers(s.createCommandBlockHandler()), + interp.ExecHandlers(s.blockHandler()), ) if err != nil { return "", "", fmt.Errorf("could not run command: %w", err)