From 576cebc9282cfbe57d45321105d72cc61597ce9b Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Sun, 27 Jul 2025 02:00:26 -0400 Subject: feat: Add Shell Command Execution to Custom Commands (#4917) --- .../prompt-processors/shellProcessor.test.ts | 300 +++++++++++++++++++++ .../services/prompt-processors/shellProcessor.ts | 106 ++++++++ .../cli/src/services/prompt-processors/types.ts | 5 + 3 files changed, 411 insertions(+) create mode 100644 packages/cli/src/services/prompt-processors/shellProcessor.test.ts create mode 100644 packages/cli/src/services/prompt-processors/shellProcessor.ts (limited to 'packages/cli/src/services/prompt-processors') diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts new file mode 100644 index 00000000..a2883923 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts @@ -0,0 +1,300 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { ConfirmationRequiredError, ShellProcessor } from './shellProcessor.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { CommandContext } from '../../ui/commands/types.js'; +import { Config } from '@google/gemini-cli-core'; + +const mockCheckCommandPermissions = vi.hoisted(() => vi.fn()); +const mockShellExecute = vi.hoisted(() => vi.fn()); + +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const original = await importOriginal(); + return { + ...original, + checkCommandPermissions: mockCheckCommandPermissions, + ShellExecutionService: { + execute: mockShellExecute, + }, + }; +}); + +describe('ShellProcessor', () => { + let context: CommandContext; + let mockConfig: Partial; + + beforeEach(() => { + vi.clearAllMocks(); + + mockConfig = { + getTargetDir: vi.fn().mockReturnValue('/test/dir'), + }; + + context = createMockCommandContext({ + services: { + config: mockConfig as Config, + }, + session: { + sessionShellAllowlist: new Set(), + }, + }); + + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + output: 'default shell output', + }), + }); + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + }); + + it('should not change the prompt if no shell injections are present', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'This is a simple prompt with no injections.'; + const result = await processor.process(prompt, context); + expect(result).toBe(prompt); + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should process a single valid shell injection if allowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'The current status is: !{git status}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ output: 'On branch main' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'git status', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + 'git status', + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('The current status is: On branch main'); + }); + + it('should process multiple valid shell injections if all are allowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{git status} in !{pwd}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + + mockShellExecute + .mockReturnValueOnce({ + result: Promise.resolve({ output: 'On branch main' }), + }) + .mockReturnValueOnce({ + result: Promise.resolve({ output: '/usr/home' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledTimes(2); + expect(mockShellExecute).toHaveBeenCalledTimes(2); + expect(result).toBe('On branch main in /usr/home'); + }); + + it('should throw ConfirmationRequiredError if a command is not allowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Do something dangerous: !{rm -rf /}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: ['rm -rf /'], + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + ConfirmationRequiredError, + ); + }); + + it('should throw ConfirmationRequiredError with the correct command', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Do something dangerous: !{rm -rf /}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: ['rm -rf /'], + }); + + try { + await processor.process(prompt, context); + // Fail if it doesn't throw + expect(true).toBe(false); + } catch (e) { + expect(e).toBeInstanceOf(ConfirmationRequiredError); + if (e instanceof ConfirmationRequiredError) { + expect(e.commandsToConfirm).toEqual(['rm -rf /']); + } + } + + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should throw ConfirmationRequiredError with multiple commands if multiple are disallowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd1} and !{cmd2}'; + mockCheckCommandPermissions.mockImplementation((cmd) => { + if (cmd === 'cmd1') { + return { allAllowed: false, disallowedCommands: ['cmd1'] }; + } + if (cmd === 'cmd2') { + return { allAllowed: false, disallowedCommands: ['cmd2'] }; + } + return { allAllowed: true, disallowedCommands: [] }; + }); + + try { + await processor.process(prompt, context); + // Fail if it doesn't throw + expect(true).toBe(false); + } catch (e) { + expect(e).toBeInstanceOf(ConfirmationRequiredError); + if (e instanceof ConfirmationRequiredError) { + expect(e.commandsToConfirm).toEqual(['cmd1', 'cmd2']); + } + } + }); + + it('should not execute any commands if at least one requires confirmation', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'First: !{echo "hello"}, Second: !{rm -rf /}'; + + mockCheckCommandPermissions.mockImplementation((cmd) => { + if (cmd.includes('rm')) { + return { allAllowed: false, disallowedCommands: [cmd] }; + } + return { allAllowed: true, disallowedCommands: [] }; + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + ConfirmationRequiredError, + ); + + // Ensure no commands were executed because the pipeline was halted. + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should only request confirmation for disallowed commands in a mixed prompt', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Allowed: !{ls -l}, Disallowed: !{rm -rf /}'; + + mockCheckCommandPermissions.mockImplementation((cmd) => ({ + allAllowed: !cmd.includes('rm'), + disallowedCommands: cmd.includes('rm') ? [cmd] : [], + })); + + try { + await processor.process(prompt, context); + expect.fail('Should have thrown ConfirmationRequiredError'); + } catch (e) { + expect(e).toBeInstanceOf(ConfirmationRequiredError); + if (e instanceof ConfirmationRequiredError) { + expect(e.commandsToConfirm).toEqual(['rm -rf /']); + } + } + }); + + it('should execute all commands if they are on the session allowlist', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Run !{cmd1} and !{cmd2}'; + + // Add commands to the session allowlist + context.session.sessionShellAllowlist = new Set(['cmd1', 'cmd2']); + + // checkCommandPermissions should now pass for these + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + + mockShellExecute + .mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) }) + .mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'cmd1', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'cmd2', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledTimes(2); + expect(result).toBe('Run output1 and output2'); + }); + + it('should trim whitespace from the command inside the injection', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Files: !{ ls -l }'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ output: 'total 0' }), + }); + + await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'ls -l', // Verifies that the command was trimmed + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + 'ls -l', + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + }); + + it('should handle an empty command inside the injection gracefully', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'This is weird: !{}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ output: 'empty output' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + '', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + '', + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('This is weird: empty output'); + }); +}); diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.ts b/packages/cli/src/services/prompt-processors/shellProcessor.ts new file mode 100644 index 00000000..bf811d66 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/shellProcessor.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + checkCommandPermissions, + ShellExecutionService, +} from '@google/gemini-cli-core'; + +import { CommandContext } from '../../ui/commands/types.js'; +import { IPromptProcessor } from './types.js'; + +export class ConfirmationRequiredError extends Error { + constructor( + message: string, + public commandsToConfirm: string[], + ) { + super(message); + this.name = 'ConfirmationRequiredError'; + } +} + +/** + * Finds all instances of shell command injections (`!{...}`) in a prompt, + * executes them, and replaces the injection site with the command's output. + * + * This processor ensures that only allowlisted commands are executed. If a + * disallowed command is found, it halts execution and reports an error. + */ +export class ShellProcessor implements IPromptProcessor { + /** + * A regular expression to find all instances of `!{...}`. The inner + * capture group extracts the command itself. + */ + private static readonly SHELL_INJECTION_REGEX = /!\{([^}]*)\}/g; + + /** + * @param commandName The name of the custom command being executed, used + * for logging and error messages. + */ + constructor(private readonly commandName: string) {} + + async process(prompt: string, context: CommandContext): Promise { + const { config, sessionShellAllowlist } = { + ...context.services, + ...context.session, + }; + const commandsToExecute: Array<{ fullMatch: string; command: string }> = []; + const commandsToConfirm = new Set(); + + const matches = [...prompt.matchAll(ShellProcessor.SHELL_INJECTION_REGEX)]; + if (matches.length === 0) { + return prompt; // No shell commands, nothing to do. + } + + // Discover all commands and check permissions. + for (const match of matches) { + const command = match[1].trim(); + const { allAllowed, disallowedCommands, blockReason, isHardDenial } = + checkCommandPermissions(command, config!, sessionShellAllowlist); + + if (!allAllowed) { + // If it's a hard denial, this is a non-recoverable security error. + if (isHardDenial) { + throw new Error( + `${this.commandName} cannot be run. ${blockReason || 'A shell command in this custom command is explicitly blocked in your config settings.'}`, + ); + } + + // Add each soft denial disallowed command to the set for confirmation. + disallowedCommands.forEach((uc) => commandsToConfirm.add(uc)); + } + commandsToExecute.push({ fullMatch: match[0], command }); + } + + // If any commands require confirmation, throw a special error to halt the + // pipeline and trigger the UI flow. + if (commandsToConfirm.size > 0) { + throw new ConfirmationRequiredError( + 'Shell command confirmation required', + Array.from(commandsToConfirm), + ); + } + + // Execute all commands (only runs if no confirmation was needed). + let processedPrompt = prompt; + for (const { fullMatch, command } of commandsToExecute) { + const { result } = ShellExecutionService.execute( + command, + config!.getTargetDir(), + () => {}, // No streaming needed. + new AbortController().signal, // For now, we don't support cancellation from here. + ); + + const executionResult = await result; + processedPrompt = processedPrompt.replace( + fullMatch, + executionResult.output, + ); + } + + return processedPrompt; + } +} diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts index 2ca61062..2653d2b7 100644 --- a/packages/cli/src/services/prompt-processors/types.ts +++ b/packages/cli/src/services/prompt-processors/types.ts @@ -35,3 +35,8 @@ export interface IPromptProcessor { * The placeholder string for shorthand argument injection in custom commands. */ export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}'; + +/** + * The trigger string for shell command injection in custom commands. + */ +export const SHELL_INJECTION_TRIGGER = '!{'; -- cgit v1.2.3