Skip to content

Commit 7a47e33

Browse files
authored
feat(chat): Add warning UI for high risk bash command (aws#6895)
## Problem AppSec has required that we give the user an indication if they are about to run a potentially dangerous command ## Solution - Parse bash command to determine whether or not commands are dangerous or not. - Display message warning to user ## Tests ![Screenshot 2025-03-31 at 1 03 14 PM](https://github.com/user-attachments/assets/a3bcd519-46de-47c9-a947-32f58dd6ca51) --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 44a7dbc commit 7a47e33

File tree

9 files changed

+210
-74
lines changed

9 files changed

+210
-74
lines changed

package-lock.json

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/core/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -492,7 +492,8 @@
492492
"umd-compat-loader": "^2.1.2",
493493
"vue-loader": "^17.2.2",
494494
"vue-style-loader": "^4.1.3",
495-
"webfont": "^11.2.26"
495+
"webfont": "^11.2.26",
496+
"shlex": "^2.1.2"
496497
},
497498
"dependencies": {
498499
"@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",

packages/core/src/codewhispererChat/controllers/chat/controller.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -662,7 +662,9 @@ export class ChatController {
662662
try {
663663
await ToolUtils.validate(tool)
664664

665-
const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse)
665+
const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
666+
requiresAcceptance: false,
667+
})
666668
const output = await ToolUtils.invoke(tool, chatStream)
667669

668670
toolResults.push({

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import { ToolType, ToolUtils } from '../../../tools/toolUtils'
4545
import { ChatStream } from '../../../tools/chatStream'
4646
import path from 'path'
4747
import { getWorkspaceForFile } from '../../../../shared/utilities/workspaceUtils'
48+
import { CommandValidation } from '../../../tools/executeBash'
4849

4950
export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
5051

@@ -220,11 +221,12 @@ export class Messenger {
220221
if (tool.type === ToolType.FsWrite) {
221222
session.setShowDiffOnFileWrite(true)
222223
}
223-
const requiresAcceptance = ToolUtils.requiresAcceptance(tool)
224-
const chatStream = new ChatStream(this, tabID, triggerID, toolUse, requiresAcceptance)
224+
const validation = ToolUtils.requiresAcceptance(tool)
225+
226+
const chatStream = new ChatStream(this, tabID, triggerID, toolUse, validation)
225227
ToolUtils.queueDescription(tool, chatStream)
226228

227-
if (!requiresAcceptance) {
229+
if (!validation.requiresAcceptance) {
228230
// 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.
229231
this.dispatcher.sendCustomFormActionMessage(
230232
new CustomFormActionMessage(tabID, {
@@ -432,17 +434,21 @@ export class Messenger {
432434
tabID: string,
433435
triggerID: string,
434436
toolUse: ToolUse | undefined,
435-
requiresAcceptance = false
437+
validation: CommandValidation
436438
) {
437439
const buttons: ChatItemButton[] = []
438440
let fileList: ChatItemContent['fileList'] = undefined
439-
if (requiresAcceptance && toolUse?.name === ToolType.ExecuteBash) {
441+
if (validation.requiresAcceptance && toolUse?.name === ToolType.ExecuteBash) {
440442
buttons.push({
441443
id: 'confirm-tool-use',
442444
text: 'Confirm',
443445
position: 'outside',
444446
status: 'info',
445447
})
448+
449+
if (validation.warning) {
450+
message = validation.warning + message
451+
}
446452
} else if (toolUse?.name === ToolType.FsWrite) {
447453
// FileList
448454
const absoluteFilePath = (toolUse?.input as any).path
@@ -471,7 +477,7 @@ export class Messenger {
471477
this.dispatcher.sendChatMessage(
472478
new ChatMessage(
473479
{
474-
message,
480+
message: message,
475481
messageType: 'answer-part',
476482
followUps: undefined,
477483
followUpsHeader: undefined,

packages/core/src/codewhispererChat/tools/chatStream.ts

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { Writable } from 'stream'
77
import { getLogger } from '../../shared/logger/logger'
88
import { Messenger } from '../controllers/chat/messenger/messenger'
99
import { ToolUse } from '@amzn/codewhisperer-streaming'
10+
import { CommandValidation } from './executeBash'
1011

1112
/**
1213
* A writable stream that feeds each chunk/line to the chat UI.
@@ -20,7 +21,7 @@ export class ChatStream extends Writable {
2021
private readonly tabID: string,
2122
private readonly triggerID: string,
2223
private readonly toolUse: ToolUse | undefined,
23-
private readonly requiresAcceptance = false,
24+
private readonly validation: CommandValidation,
2425
private readonly logger = getLogger('chatStream')
2526
) {
2627
super()
@@ -37,7 +38,7 @@ export class ChatStream extends Writable {
3738
this.tabID,
3839
this.triggerID,
3940
this.toolUse,
40-
this.requiresAcceptance
41+
this.validation
4142
)
4243
callback()
4344
}
@@ -49,7 +50,7 @@ export class ChatStream extends Writable {
4950
this.tabID,
5051
this.triggerID,
5152
this.toolUse,
52-
this.requiresAcceptance
53+
this.validation
5354
)
5455
}
5556
callback()

packages/core/src/codewhispererChat/tools/executeBash.ts

Lines changed: 167 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,130 @@
55

66
import { Writable } from 'stream'
77
import { getLogger } from '../../shared/logger/logger'
8-
import { fs } from '../../shared/fs/fs' // e.g. for getUserHomeDir()
8+
import { fs } from '../../shared/fs/fs'
99
import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils'
1010
import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
11+
import { split } from 'shlex'
1112

12-
export const readOnlyCommands: string[] = ['ls', 'cat', 'echo', 'pwd', 'which', 'head', 'tail']
13+
export enum CommandCategory {
14+
ReadOnly,
15+
HighRisk,
16+
Destructive,
17+
}
18+
19+
export const dangerousPatterns = new Set(['<(', '$(', '`', '>', '&&', '||'])
20+
export const commandCategories = new Map<string, CommandCategory>([
21+
// ReadOnly commands
22+
['ls', CommandCategory.ReadOnly],
23+
['cat', CommandCategory.ReadOnly],
24+
['bat', CommandCategory.ReadOnly],
25+
['pwd', CommandCategory.ReadOnly],
26+
['echo', CommandCategory.ReadOnly],
27+
['file', CommandCategory.ReadOnly],
28+
['less', CommandCategory.ReadOnly],
29+
['more', CommandCategory.ReadOnly],
30+
['tree', CommandCategory.ReadOnly],
31+
['find', CommandCategory.ReadOnly],
32+
['top', CommandCategory.ReadOnly],
33+
['htop', CommandCategory.ReadOnly],
34+
['ps', CommandCategory.ReadOnly],
35+
['df', CommandCategory.ReadOnly],
36+
['du', CommandCategory.ReadOnly],
37+
['free', CommandCategory.ReadOnly],
38+
['uname', CommandCategory.ReadOnly],
39+
['date', CommandCategory.ReadOnly],
40+
['whoami', CommandCategory.ReadOnly],
41+
['which', CommandCategory.ReadOnly],
42+
['ping', CommandCategory.ReadOnly],
43+
['ifconfig', CommandCategory.ReadOnly],
44+
['ip', CommandCategory.ReadOnly],
45+
['netstat', CommandCategory.ReadOnly],
46+
['ss', CommandCategory.ReadOnly],
47+
['dig', CommandCategory.ReadOnly],
48+
['grep', CommandCategory.ReadOnly],
49+
['wc', CommandCategory.ReadOnly],
50+
['sort', CommandCategory.ReadOnly],
51+
['diff', CommandCategory.ReadOnly],
52+
['head', CommandCategory.ReadOnly],
53+
['tail', CommandCategory.ReadOnly],
54+
55+
// HighRisk commands
56+
['chmod', CommandCategory.HighRisk],
57+
['chown', CommandCategory.HighRisk],
58+
['mv', CommandCategory.HighRisk],
59+
['cp', CommandCategory.HighRisk],
60+
['ln', CommandCategory.HighRisk],
61+
['mount', CommandCategory.HighRisk],
62+
['umount', CommandCategory.HighRisk],
63+
['kill', CommandCategory.HighRisk],
64+
['killall', CommandCategory.HighRisk],
65+
['pkill', CommandCategory.HighRisk],
66+
['iptables', CommandCategory.HighRisk],
67+
['route', CommandCategory.HighRisk],
68+
['systemctl', CommandCategory.HighRisk],
69+
['service', CommandCategory.HighRisk],
70+
['crontab', CommandCategory.HighRisk],
71+
['at', CommandCategory.HighRisk],
72+
['tar', CommandCategory.HighRisk],
73+
['awk', CommandCategory.HighRisk],
74+
['sed', CommandCategory.HighRisk],
75+
['wget', CommandCategory.HighRisk],
76+
['curl', CommandCategory.HighRisk],
77+
['nc', CommandCategory.HighRisk],
78+
['ssh', CommandCategory.HighRisk],
79+
['scp', CommandCategory.HighRisk],
80+
['ftp', CommandCategory.HighRisk],
81+
['sftp', CommandCategory.HighRisk],
82+
['rsync', CommandCategory.HighRisk],
83+
['chroot', CommandCategory.HighRisk],
84+
['lsof', CommandCategory.HighRisk],
85+
['strace', CommandCategory.HighRisk],
86+
['gdb', CommandCategory.HighRisk],
87+
88+
// Destructive commands
89+
['rm', CommandCategory.Destructive],
90+
['dd', CommandCategory.Destructive],
91+
['mkfs', CommandCategory.Destructive],
92+
['fdisk', CommandCategory.Destructive],
93+
['shutdown', CommandCategory.Destructive],
94+
['reboot', CommandCategory.Destructive],
95+
['poweroff', CommandCategory.Destructive],
96+
['sudo', CommandCategory.Destructive],
97+
['su', CommandCategory.Destructive],
98+
['useradd', CommandCategory.Destructive],
99+
['userdel', CommandCategory.Destructive],
100+
['passwd', CommandCategory.Destructive],
101+
['visudo', CommandCategory.Destructive],
102+
['insmod', CommandCategory.Destructive],
103+
['rmmod', CommandCategory.Destructive],
104+
['modprobe', CommandCategory.Destructive],
105+
['apt', CommandCategory.Destructive],
106+
['yum', CommandCategory.Destructive],
107+
['dnf', CommandCategory.Destructive],
108+
['pacman', CommandCategory.Destructive],
109+
['perl', CommandCategory.Destructive],
110+
['python', CommandCategory.Destructive],
111+
['bash', CommandCategory.Destructive],
112+
['sh', CommandCategory.Destructive],
113+
['exec', CommandCategory.Destructive],
114+
['eval', CommandCategory.Destructive],
115+
['xargs', CommandCategory.Destructive],
116+
])
13117
export const maxBashToolResponseSize: number = 1024 * 1024 // 1MB
14118
export const lineCount: number = 1024
15-
export const dangerousPatterns: string[] = ['|', '<(', '$(', '`', '>', '&&', '||']
119+
export const destructiveCommandWarningMessage = '⚠️ WARNING: Destructive command detected:\n\n'
120+
export const highRiskCommandWarningMessage = '⚠️ WARNING: High risk command detected:\n\n'
16121

17122
export interface ExecuteBashParams {
18123
command: string
19124
cwd?: string
20125
}
21126

127+
export interface CommandValidation {
128+
requiresAcceptance: boolean
129+
warning?: string
130+
}
131+
22132
export class ExecuteBash {
23133
private readonly command: string
24134
private readonly workingDirectory?: string
@@ -34,7 +144,7 @@ export class ExecuteBash {
34144
throw new Error('Bash command cannot be empty.')
35145
}
36146

37-
const args = ExecuteBash.parseCommand(this.command)
147+
const args = split(this.command)
38148
if (!args || args.length === 0) {
39149
throw new Error('No command found.')
40150
}
@@ -46,22 +156,67 @@ export class ExecuteBash {
46156
}
47157
}
48158

49-
public requiresAcceptance(): boolean {
159+
public requiresAcceptance(): CommandValidation {
50160
try {
51-
const args = ExecuteBash.parseCommand(this.command)
161+
const args = split(this.command)
52162
if (!args || args.length === 0) {
53-
return true
163+
return { requiresAcceptance: true }
54164
}
55165

56-
if (args.some((arg) => dangerousPatterns.some((pattern) => arg.includes(pattern)))) {
57-
return true
166+
// Split commands by pipe and process each segment
167+
let currentCmd: string[] = []
168+
const allCommands: string[][] = []
169+
170+
for (const arg of args) {
171+
if (arg === '|') {
172+
if (currentCmd.length > 0) {
173+
allCommands.push(currentCmd)
174+
}
175+
currentCmd = []
176+
} else if (arg.includes('|')) {
177+
return { requiresAcceptance: true }
178+
} else {
179+
currentCmd.push(arg)
180+
}
181+
}
182+
183+
if (currentCmd.length > 0) {
184+
allCommands.push(currentCmd)
58185
}
59186

60-
const command = args[0]
61-
return !readOnlyCommands.includes(command)
187+
for (const cmdArgs of allCommands) {
188+
if (cmdArgs.length === 0) {
189+
return { requiresAcceptance: true }
190+
}
191+
192+
const command = cmdArgs[0]
193+
const category = commandCategories.get(command)
194+
195+
switch (category) {
196+
case CommandCategory.Destructive:
197+
return { requiresAcceptance: true, warning: destructiveCommandWarningMessage }
198+
case CommandCategory.HighRisk:
199+
return {
200+
requiresAcceptance: true,
201+
warning: highRiskCommandWarningMessage,
202+
}
203+
case CommandCategory.ReadOnly:
204+
if (
205+
cmdArgs.some((arg) =>
206+
Array.from(dangerousPatterns).some((pattern) => arg.includes(pattern))
207+
)
208+
) {
209+
return { requiresAcceptance: true, warning: highRiskCommandWarningMessage }
210+
}
211+
return { requiresAcceptance: false }
212+
default:
213+
return { requiresAcceptance: true, warning: highRiskCommandWarningMessage }
214+
}
215+
}
216+
return { requiresAcceptance: true }
62217
} catch (error) {
63218
this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`)
64-
return true
219+
return { requiresAcceptance: true }
65220
}
66221
}
67222

@@ -167,43 +322,6 @@ export class ExecuteBash {
167322
return output
168323
}
169324

170-
private static parseCommand(command: string): string[] | undefined {
171-
const result: string[] = []
172-
let current = ''
173-
let inQuote: string | undefined
174-
let escaped = false
175-
176-
for (const char of command) {
177-
if (escaped) {
178-
current += char
179-
escaped = false
180-
} else if (char === '\\') {
181-
escaped = true
182-
} else if (inQuote) {
183-
if (char === inQuote) {
184-
inQuote = undefined
185-
} else {
186-
current += char
187-
}
188-
} else if (char === '"' || char === "'") {
189-
inQuote = char
190-
} else if (char === ' ' || char === '\t') {
191-
if (current) {
192-
result.push(current)
193-
current = ''
194-
}
195-
} else {
196-
current += char
197-
}
198-
}
199-
200-
if (current) {
201-
result.push(current)
202-
}
203-
204-
return result
205-
}
206-
207325
public queueDescription(updates: Writable): void {
208326
updates.write(`I will run the following shell command:\n`)
209327
updates.write('```bash\n' + this.command + '\n```')

0 commit comments

Comments
 (0)