feat(chat): Add warning UI for high risk bash command

tsmithsz · tsmithsz · commit f7031f2c7504 · 2025-03-31T16:25:07.000-07:00
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/core/package.json b/packages/core/package.json
@@ -492,7 +492,8 @@
         "umd-compat-loader": "^2.1.2",
         "vue-loader": "^17.2.2",
         "vue-style-loader": "^4.1.3",
-        "webfont": "^11.2.26"
+        "webfont": "^11.2.26",
+        "shlex": "^2.1.2"
     },
     "dependencies": {
         "@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -662,7 +662,10 @@ export class ChatController {
                     try {
                         await ToolUtils.validate(tool)
 
-                        const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse)
+                        const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
+                            requiresAcceptance: false,
+                            isHighRisk: false,
+                        })
                         const output = await ToolUtils.invoke(tool, chatStream)
 
                         toolResults.push({
diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
@@ -45,6 +45,7 @@ import { ToolType, ToolUtils } from '../../../tools/toolUtils'
 import { ChatStream } from '../../../tools/chatStream'
 import path from 'path'
 import { getWorkspaceForFile } from '../../../../shared/utilities/workspaceUtils'
+import { CommandValidation } from '../../../tools/executeBash'
 
 export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
 
@@ -220,11 +221,12 @@ export class Messenger {
                             if (tool.type === ToolType.FsWrite) {
                                 session.setShowDiffOnFileWrite(true)
                             }
-                            const requiresAcceptance = ToolUtils.requiresAcceptance(tool)
-                            const chatStream = new ChatStream(this, tabID, triggerID, toolUse, requiresAcceptance)
+                            const validation = ToolUtils.requiresAcceptance(tool)
+
+                            const chatStream = new ChatStream(this, tabID, triggerID, toolUse, validation)
                             ToolUtils.queueDescription(tool, chatStream)
 
-                            if (!requiresAcceptance) {
+                            if (!validation.requiresAcceptance) {
                                 // Need separate id for read tool and safe bash command execution as 'confirm-tool-use' id is required to change button status from `Confirm` to `Confirmed` state in cwChatConnector.ts which will impact generic tool execution.
                                 this.dispatcher.sendCustomFormActionMessage(
                                     new CustomFormActionMessage(tabID, {
@@ -432,17 +434,21 @@ export class Messenger {
         tabID: string,
         triggerID: string,
         toolUse: ToolUse | undefined,
-        requiresAcceptance = false
+        validation: CommandValidation
     ) {
         const buttons: ChatItemButton[] = []
         let fileList: ChatItemContent['fileList'] = undefined
-        if (requiresAcceptance && toolUse?.name === ToolType.ExecuteBash) {
+        if (validation.requiresAcceptance && toolUse?.name === ToolType.ExecuteBash) {
             buttons.push({
                 id: 'confirm-tool-use',
                 text: 'Confirm',
                 position: 'outside',
                 status: 'info',
             })
+
+            if (validation.isHighRisk) {
+                message = '⚠️ WARNING: High risk command detected:\n\n' + message
+            }
         } else if (toolUse?.name === ToolType.FsWrite) {
             // FileList
             const absoluteFilePath = (toolUse?.input as any).path
@@ -471,7 +477,7 @@ export class Messenger {
         this.dispatcher.sendChatMessage(
             new ChatMessage(
                 {
-                    message,
+                    message: message,
                     messageType: 'answer-part',
                     followUps: undefined,
                     followUpsHeader: undefined,
diff --git a/packages/core/src/codewhispererChat/tools/chatStream.ts b/packages/core/src/codewhispererChat/tools/chatStream.ts
@@ -7,6 +7,7 @@ import { Writable } from 'stream'
 import { getLogger } from '../../shared/logger/logger'
 import { Messenger } from '../controllers/chat/messenger/messenger'
 import { ToolUse } from '@amzn/codewhisperer-streaming'
+import { CommandValidation } from './executeBash'
 
 /**
  * A writable stream that feeds each chunk/line to the chat UI.
@@ -20,7 +21,7 @@ export class ChatStream extends Writable {
         private readonly tabID: string,
         private readonly triggerID: string,
         private readonly toolUse: ToolUse | undefined,
-        private readonly requiresAcceptance = false,
+        private readonly validation: CommandValidation,
         private readonly logger = getLogger('chatStream')
     ) {
         super()
@@ -37,7 +38,7 @@ export class ChatStream extends Writable {
             this.tabID,
             this.triggerID,
             this.toolUse,
-            this.requiresAcceptance
+            this.validation
         )
         callback()
     }
@@ -49,7 +50,7 @@ export class ChatStream extends Writable {
                 this.tabID,
                 this.triggerID,
                 this.toolUse,
-                this.requiresAcceptance
+                this.validation
             )
         }
         callback()
diff --git a/packages/core/src/codewhispererChat/tools/executeBash.ts b/packages/core/src/codewhispererChat/tools/executeBash.ts
@@ -5,20 +5,127 @@
 
 import { Writable } from 'stream'
 import { getLogger } from '../../shared/logger/logger'
-import { fs } from '../../shared/fs/fs' // e.g. for getUserHomeDir()
+import { fs } from '../../shared/fs/fs'
 import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils'
 import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { split } from 'shlex'
 
-export const readOnlyCommands: string[] = ['ls', 'cat', 'echo', 'pwd', 'which', 'head', 'tail']
+export const readOnlyCommands = new Set<string>([
+    'ls',
+    'cat',
+    'bat',
+    'pwd',
+    'echo',
+    'file',
+    'less',
+    'more',
+    'tree',
+    'find',
+    'top',
+    'htop',
+    'ps',
+    'df',
+    'du',
+    'free',
+    'uname',
+    'date',
+    'whoami',
+    'which',
+    'ping',
+    'ifconfig',
+    'ip',
+    'netstat',
+    'ss',
+    'dig',
+    'grep',
+    'wc',
+    'sort',
+    'diff',
+    'head',
+    'tail',
+])
+export const dangerousPatterns = new Set(['<(', '$(', '`', '>', '&&', '||'])
+const highRiskCommands = new Map<string, string>([
+    ['rm', 'This command can delete files and directories'],
+    ['chmod', 'This command can modify file permissions'],
+    ['chown', 'This command can change file ownership'],
+    ['dd', 'This command can overwrite disk data'],
+    ['sudo', 'This command can execute with elevated privileges'],
+    ['su', 'This command can switch user accounts'],
+    ['mv', 'This command can move or rename files'],
+    ['cp', 'This command can copy files, potentially overwriting existing ones'],
+    ['ln', 'This command can create links, potentially to sensitive files'],
+    ['mkfs', 'This command can format filesystems'],
+    ['fdisk', 'This command can modify disk partitions'],
+    ['mount', 'This command can mount filesystems'],
+    ['umount', 'This command can unmount filesystems'],
+    ['shutdown', 'This command can shut down the system'],
+    ['reboot', 'This command can restart the system'],
+    ['poweroff', 'This command can power off the system'],
+    ['kill', 'This command can terminate processes'],
+    ['killall', 'This command can terminate multiple processes'],
+    ['pkill', 'This command can terminate processes by name'],
+    ['iptables', 'This command can modify firewall rules'],
+    ['ifconfig', 'This command can configure network interfaces'],
+    ['route', 'This command can modify network routing tables'],
+    ['useradd', 'This command can create new user accounts'],
+    ['userdel', 'This command can delete user accounts'],
+    ['passwd', 'This command can change user passwords'],
+    ['visudo', 'This command can modify sudo privileges'],
+    ['crontab', 'This command can schedule tasks to run automatically'],
+    ['at', 'This command can schedule one-time tasks'],
+    ['systemctl', 'This command can control system services'],
+    ['service', 'This command can start, stop, or restart system services'],
+    ['insmod', 'This command can insert a module into the Linux kernel'],
+    ['rmmod', 'This command can remove a module from the Linux kernel'],
+    ['modprobe', 'This command can add or remove modules from the Linux kernel'],
+    ['apt', 'This command can install, update, or remove software packages'],
+    ['yum', 'This command can install, update, or remove software packages'],
+    ['dnf', 'This command can install, update, or remove software packages'],
+    ['pacman', 'This command can install, update, or remove software packages'],
+    ['tar', 'This command can create or extract archives, potentially overwriting files'],
+    ['find', 'This command can search for files and execute commands on them'],
+    ['grep', 'This command can search file contents and potentially reveal sensitive information'],
+    ['awk', 'This command can process and modify text files'],
+    ['sed', 'This command can modify file contents'],
+    ['perl', 'This command can execute Perl scripts'],
+    ['python', 'This command can execute Python scripts'],
+    ['bash', 'This command can execute bash scripts'],
+    ['sh', 'This command can execute shell scripts'],
+    ['exec', 'This command can execute commands replacing the current process'],
+    ['eval', 'This command can execute arguments as a shell command'],
+    ['xargs', 'This command can build and execute command lines from standard input'],
+    ['nohup', 'This command can run commands immune to hangups'],
+    ['wget', 'This command can download files from the web'],
+    ['curl', 'This command can transfer data from or to a server'],
+    ['nc', 'This command can arbitrarily read and write network connections'],
+    ['ssh', 'This command can remotely access other systems'],
+    ['scp', 'This command can copy files between hosts on a network'],
+    ['ftp', 'This command can transfer files to and from a remote network'],
+    ['sftp', 'This command can securely transfer files over SSH'],
+    ['rsync', 'This command can synchronize files and directories'],
+    ['chroot', 'This command can change the root directory'],
+    ['mount', 'This command can mount file systems'],
+    ['umount', 'This command can unmount file systems'],
+    ['lsof', 'This command can list open files and the processes that opened them'],
+    ['strace', 'This command can trace system calls and signals'],
+    ['gdb', 'This command can debug programs'],
+    ['strings', 'This command can print printable characters in files'],
+])
 export const maxBashToolResponseSize: number = 1024 * 1024 // 1MB
 export const lineCount: number = 1024
-export const dangerousPatterns: string[] = ['|', '<(', '$(', '`', '>', '&&', '||']
 
 export interface ExecuteBashParams {
     command: string
     cwd?: string
 }
 
+export interface CommandValidation {
+    requiresAcceptance: boolean
+    isHighRisk: boolean
+    reason?: string
+}
+
 export class ExecuteBash {
     private readonly command: string
     private readonly workingDirectory?: string
@@ -34,7 +141,7 @@ export class ExecuteBash {
             throw new Error('Bash command cannot be empty.')
         }
 
-        const args = ExecuteBash.parseCommand(this.command)
+        const args = split(this.command)
         if (!args || args.length === 0) {
             throw new Error('No command found.')
         }
@@ -46,22 +153,60 @@ export class ExecuteBash {
         }
     }
 
-    public requiresAcceptance(): boolean {
+    public requiresAcceptance(): CommandValidation {
         try {
-            const args = ExecuteBash.parseCommand(this.command)
+            const args = split(this.command)
             if (!args || args.length === 0) {
-                return true
+                return { requiresAcceptance: true, isHighRisk: false }
+            }
+
+            // Split commands by pipe and process each segment
+            let currentCmd: string[] = []
+            const allCommands: string[][] = []
+
+            for (const arg of args) {
+                if (arg === '|') {
+                    if (currentCmd.length > 0) {
+                        allCommands.push(currentCmd)
+                    }
+                    currentCmd = []
+                } else if (arg.includes('|')) {
+                    return { requiresAcceptance: true, isHighRisk: false }
+                } else {
+                    currentCmd.push(arg)
+                }
             }
 
-            if (args.some((arg) => dangerousPatterns.some((pattern) => arg.includes(pattern)))) {
-                return true
+            if (currentCmd.length > 0) {
+                allCommands.push(currentCmd)
             }
 
-            const command = args[0]
-            return !readOnlyCommands.includes(command)
+            for (const cmdArgs of allCommands) {
+                if (cmdArgs.length === 0) {
+                    return { requiresAcceptance: true, isHighRisk: false }
+                }
+
+                if (cmdArgs.some((arg) => Array.from(dangerousPatterns).some((pattern) => arg.includes(pattern)))) {
+                    return { requiresAcceptance: true, isHighRisk: false }
+                }
+
+                const command = cmdArgs[0]
+
+                if (highRiskCommands.has(command)) {
+                    const reason = highRiskCommands.get(command)
+                    this.logger.warn(`WARNING: High risk command detected: ${command}`)
+                    this.logger.warn(`Reason: ${reason}`)
+                    return { requiresAcceptance: true, isHighRisk: true, reason: reason }
+                }
+
+                if (!readOnlyCommands.has(command)) {
+                    return { requiresAcceptance: true, isHighRisk: false }
+                }
+            }
+            return { requiresAcceptance: false, isHighRisk: false }
         } catch (error) {
             this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`)
-            return true
+            return { requiresAcceptance: true, isHighRisk: false }
         }
     }
 
@@ -167,43 +312,6 @@ export class ExecuteBash {
         return output
     }
 
-    private static parseCommand(command: string): string[] | undefined {
-        const result: string[] = []
-        let current = ''
-        let inQuote: string | undefined
-        let escaped = false
-
-        for (const char of command) {
-            if (escaped) {
-                current += char
-                escaped = false
-            } else if (char === '\\') {
-                escaped = true
-            } else if (inQuote) {
-                if (char === inQuote) {
-                    inQuote = undefined
-                } else {
-                    current += char
-                }
-            } else if (char === '"' || char === "'") {
-                inQuote = char
-            } else if (char === ' ' || char === '\t') {
-                if (current) {
-                    result.push(current)
-                    current = ''
-                }
-            } else {
-                current += char
-            }
-        }
-
-        if (current) {
-            result.push(current)
-        }
-
-        return result
-    }
-
     public queueDescription(updates: Writable): void {
         updates.write(`I will run the following shell command:\n`)
         updates.write('```bash\n' + this.command + '\n```')
diff --git a/packages/core/src/codewhispererChat/tools/toolUtils.ts b/packages/core/src/codewhispererChat/tools/toolUtils.ts
@@ -5,7 +5,7 @@
 import { Writable } from 'stream'
 import { FsRead, FsReadParams } from './fsRead'
 import { FsWrite, FsWriteParams } from './fsWrite'
-import { ExecuteBash, ExecuteBashParams } from './executeBash'
+import { CommandValidation, ExecuteBash, ExecuteBashParams } from './executeBash'
 import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
 import { InvokeOutput } from './toolShared'
 import { ListDirectory, ListDirectoryParams } from './listDirectory'
@@ -37,16 +37,16 @@ export class ToolUtils {
         }
     }
 
-    static requiresAcceptance(tool: Tool) {
+    static requiresAcceptance(tool: Tool): CommandValidation {
         switch (tool.type) {
             case ToolType.FsRead:
-                return false
+                return { requiresAcceptance: false, isHighRisk: false }
             case ToolType.FsWrite:
-                return true
+                return { requiresAcceptance: true, isHighRisk: false }
             case ToolType.ExecuteBash:
                 return tool.tool.requiresAcceptance()
             case ToolType.ListDirectory:
-                return false
+                return { requiresAcceptance: false, isHighRisk: false }
         }
     }
 
diff --git a/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts b/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts
diff --git a/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts b/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts
