Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 208 additions & 0 deletions packages/core/src/codewhispererChat/tools/executeBash.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,208 @@
/*!
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

import { Writable } from 'stream'
import { getLogger } from '../../shared/logger/logger'
import { fs } from '../../shared/fs/fs' // e.g. for getUserHomeDir()
import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils'
import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'

export const readOnlyCommands: string[] = ['ls', 'cat', 'echo', 'pwd', 'which', 'head', 'tail']
export const maxBashToolResponseSize: number = 1024 * 1024 // 1MB
export const lineCount: number = 1024
export const dangerousPatterns: string[] = ['|', '<(', '$(', '`', '>', '&&', '||']

export interface ExecuteBashParams {
command: string
cwd?: string
}

export class ExecuteBash {
private readonly command: string
private readonly workingDirectory?: string
private readonly logger = getLogger('executeBash')

constructor(params: ExecuteBashParams) {
this.command = params.command
this.workingDirectory = params.cwd ? sanitizePath(params.cwd) : fs.getUserHomeDir()
}

public async validate(): Promise<void> {
if (!this.command.trim()) {
throw new Error('Bash command cannot be empty.')
}

const args = ExecuteBash.parseCommand(this.command)
if (!args || args.length === 0) {
throw new Error('No command found.')
}

try {
await ExecuteBash.whichCommand(args[0])
} catch {
throw new Error(`Command "${args[0]}" not found on PATH.`)
}
}

public requiresAcceptance(): boolean {
try {
const args = ExecuteBash.parseCommand(this.command)
if (!args || args.length === 0) {
return true
}

if (args.some((arg) => dangerousPatterns.some((pattern) => arg.includes(pattern)))) {
return true
}

const command = args[0]
return !readOnlyCommands.includes(command)
} catch (error) {
this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`)
return true
}
}

public async invoke(updates: Writable): Promise<InvokeOutput> {
this.logger.info(`Invoking bash command: "${this.command}" in cwd: "${this.workingDirectory}"`)

return new Promise(async (resolve, reject) => {
this.logger.debug(`Spawning process with command: bash -c "${this.command}" (cwd=${this.workingDirectory})`)

const stdoutBuffer: string[] = []
const stderrBuffer: string[] = []

const childProcessOptions: ChildProcessOptions = {
spawnOptions: {
cwd: this.workingDirectory,
stdio: ['pipe', 'pipe', 'pipe'],
},
collect: false,
waitForStreams: true,
onStdout: (chunk: string) => {
ExecuteBash.handleChunk(chunk, stdoutBuffer, updates)
},
onStderr: (chunk: string) => {
ExecuteBash.handleChunk(chunk, stderrBuffer, updates)
},
}

const childProcess = new ChildProcess('bash', ['-c', this.command], childProcessOptions)

try {
const result = await childProcess.run()
const exitStatus = result.exitCode ?? 0
const stdout = stdoutBuffer.join('\n')
const stderr = stderrBuffer.join('\n')
const [stdoutTrunc, stdoutSuffix] = ExecuteBash.truncateSafelyWithSuffix(
stdout,
maxBashToolResponseSize / 3
)
const [stderrTrunc, stderrSuffix] = ExecuteBash.truncateSafelyWithSuffix(
stderr,
maxBashToolResponseSize / 3
)

const outputJson = {
exitStatus: exitStatus.toString(),
stdout: stdoutTrunc + (stdoutSuffix ? ' ... truncated' : ''),
stderr: stderrTrunc + (stderrSuffix ? ' ... truncated' : ''),
}

return {
output: {
kind: OutputKind.Json,
content: outputJson,
},
}
} catch (err: any) {
this.logger.error(`Failed to execute bash command '${this.command}': ${err.message}`)
throw new Error(`Failed to execute command: ${err.message}`)
}
})
}

private static handleChunk(chunk: string, buffer: string[], updates: Writable) {
const lines = chunk.split(/\r?\n/)
for (const line of lines) {
updates.write(`${line}\n`)
buffer.push(line)
if (buffer.length > lineCount) {
buffer.shift()
}
}
}

private static truncateSafelyWithSuffix(str: string, maxLength: number): [string, boolean] {
if (str.length > maxLength) {
return [str.substring(0, maxLength), true]
}
return [str, false]
}

private static async whichCommand(cmd: string): Promise<string> {
const cp = new ChildProcess('which', [cmd], {
collect: true,
waitForStreams: true,
})
const result = await cp.run()

if (result.exitCode !== 0) {
throw new Error(`Command "${cmd}" not found on PATH.`)
}

const output = result.stdout.trim()
if (!output) {
throw new Error(`Command "${cmd}" found but 'which' returned empty output.`)
}
return output
}

private static parseCommand(command: string): string[] | undefined {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parsing bash/shell commands is non-trivial. Is this intended to support arbitrary bash command, or only very simple cases?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here's an example: #7230

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: `)

if (this.command.length > 20) {
updates.write('\n')
}
updates.write(`\x1b[32m${this.command}\x1b[0m\n`)
}
}
8 changes: 4 additions & 4 deletions packages/core/src/codewhispererChat/tools/fsRead.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ export interface FsReadParams {
export class FsRead {
private fsPath: string
private readonly readRange?: number[]
private type?: boolean // true for file, false for directory
private isFile?: boolean // true for file, false for directory
private readonly logger = getLogger('fsRead')

constructor(params: FsReadParams) {
Expand Down Expand Up @@ -44,19 +44,19 @@ export class FsRead {
throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed. (${err})`)
}

this.type = await fs.existsFile(fileUri)
this.isFile = await fs.existsFile(fileUri)
this.logger.debug(`Validation succeeded for path: ${this.fsPath}`)
}

public async invoke(): Promise<InvokeOutput> {
try {
const fileUri = vscode.Uri.file(this.fsPath)

if (this.type) {
if (this.isFile) {
const fileContents = await this.readFile(fileUri)
this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`)
return this.handleFileRange(fileContents)
} else if (!this.type) {
} else if (!this.isFile) {
const maxDepth = this.getDirectoryDepth() ?? 0
const listing = await readDirectoryRecursively(fileUri, maxDepth)
return this.createOutput(listing.join('\n'))
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/shared/logger/logger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ export type LogTopic =
| 'stepfunctions'
| 'fsRead'
| 'fsWrite'
| 'executeBash'

class ErrorLog {
constructor(
Expand Down
112 changes: 112 additions & 0 deletions packages/core/src/test/codewhispererChat/tools/executeBash.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
/*!
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

import { strict as assert } from 'assert'
import sinon from 'sinon'
import { ExecuteBash } from '../../../codewhispererChat/tools/executeBash'
import { ChildProcess } from '../../../shared/utilities/processUtils'

describe('ExecuteBash Tool', () => {
let runStub: sinon.SinonStub
let invokeStub: sinon.SinonStub

beforeEach(() => {
runStub = sinon.stub(ChildProcess.prototype, 'run')
invokeStub = sinon.stub(ExecuteBash.prototype, 'invoke')
})

afterEach(() => {
sinon.restore()
})

it('pass validation for a safe command (read-only)', async () => {
runStub.resolves({
exitCode: 0,
stdout: '/bin/ls',
stderr: '',
error: undefined,
signal: undefined,
})
const execBash = new ExecuteBash({ command: 'ls' })
await execBash.validate()
})

it('fail validation if the command is empty', async () => {
const execBash = new ExecuteBash({ command: ' ' })
await assert.rejects(
execBash.validate(),
/Bash command cannot be empty/i,
'Expected an error for empty command'
)
})

it('set requiresAcceptance=true if the command has dangerous patterns', () => {
const execBash = new ExecuteBash({ command: 'ls && rm -rf /' })
const needsAcceptance = execBash.requiresAcceptance()
assert.equal(needsAcceptance, true, 'Should require acceptance for dangerous pattern')
})

it('set requiresAcceptance=false if it is a read-only command', () => {
const execBash = new ExecuteBash({ command: 'cat file.txt' })
const needsAcceptance = execBash.requiresAcceptance()
assert.equal(needsAcceptance, false, 'Read-only command should not require acceptance')
})

it('whichCommand cannot find the first arg', async () => {
runStub.resolves({
exitCode: 1,
stdout: '',
stderr: '',
error: undefined,
signal: undefined,
})

const execBash = new ExecuteBash({ command: 'noSuchCmd' })
await assert.rejects(execBash.validate(), /not found on PATH/i, 'Expected not found error from whichCommand')
})

it('whichCommand sees first arg on PATH', async () => {
runStub.resolves({
exitCode: 0,
stdout: '/usr/bin/noSuchCmd\n',
stderr: '',
error: undefined,
signal: undefined,
})

const execBash = new ExecuteBash({ command: 'noSuchCmd' })
await execBash.validate()
})

it('stub invoke() call', async () => {
invokeStub.resolves({
output: {
kind: 'json',
content: {
exitStatus: '0',
stdout: 'mocked stdout lines',
stderr: '',
},
},
})

const execBash = new ExecuteBash({ command: 'ls' })

const dummyWritable = { write: () => {} } as any
const result = await execBash.invoke(dummyWritable)

assert.strictEqual(result.output.kind, 'json')
const out = result.output.content as unknown as {
exitStatus: string
stdout: string
stderr: string
}
assert.strictEqual(out.exitStatus, '0')
assert.strictEqual(out.stdout, 'mocked stdout lines')
assert.strictEqual(out.stderr, '')

assert.strictEqual(invokeStub.callCount, 1)
})
})
Loading