106 lines
3.4 KiB
TypeScript
106 lines
3.4 KiB
TypeScript
/**
|
|
* @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<string> {
|
|
const { config, sessionShellAllowlist } = {
|
|
...context.services,
|
|
...context.session,
|
|
};
|
|
const commandsToExecute: Array<{ fullMatch: string; command: string }> = [];
|
|
const commandsToConfirm = new Set<string>();
|
|
|
|
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;
|
|
}
|
|
}
|