gemini-cli/packages/core/src/policy/shell-safety.test.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, it, expect, beforeEach } from 'vitest';
import { PolicyEngine } from './policy-engine.js';
import { PolicyDecision } from './types.js';
import type { FunctionCall } from '@google/genai';

describe('Shell Safety Policy', () => {
  let policyEngine: PolicyEngine;

  beforeEach(() => {
    policyEngine = new PolicyEngine({
      rules: [
        {
          toolName: 'run_shell_command',
          // Mimic the regex generated by toml-loader for commandPrefix = ["git log"]
          // Regex: "command":"git log(?:[\s"]|$)
          argsPattern: /"command":"git log(?:[\s"]|$)/,
          decision: PolicyDecision.ALLOW,
          priority: 1.01, // Higher priority than default
        },
      ],
      defaultDecision: PolicyDecision.ASK_USER,
    });
  });

  it('SHOULD match "git log" exactly', async () => {
    const toolCall: FunctionCall = {
      name: 'run_shell_command',
      args: { command: 'git log' },
    };
    const result = await policyEngine.check(toolCall, undefined);
    expect(result.decision).toBe(PolicyDecision.ALLOW);
  });

  it('SHOULD match "git log" with arguments', async () => {
    const toolCall: FunctionCall = {
      name: 'run_shell_command',
      args: { command: 'git log --oneline' },
    };
    const result = await policyEngine.check(toolCall, undefined);
    expect(result.decision).toBe(PolicyDecision.ALLOW);
  });

  it('SHOULD NOT match "git logout" when prefix is "git log" (strict word boundary)', async () => {
    const toolCall: FunctionCall = {
      name: 'run_shell_command',
      args: { command: 'git logout' },
    };

    // Desired behavior: Should NOT match "git log" prefix.
    // If it doesn't match, it should fall back to default decision (ASK_USER).
    const result = await policyEngine.check(toolCall, undefined);
    expect(result.decision).toBe(PolicyDecision.ASK_USER);
  });

  it('SHOULD NOT allow "git log && rm -rf /" completely when prefix is "git log" (compound command safety)', async () => {
    const toolCall: FunctionCall = {
      name: 'run_shell_command',
      args: { command: 'git log && rm -rf /' },
    };

    // Desired behavior: Should inspect all parts. "rm -rf /" is not allowed.
    // The "git log" part is ALLOW, but "rm -rf /" is ASK_USER (default).
    // Aggregate should be ASK_USER.
    const result = await policyEngine.check(toolCall, undefined);
    expect(result.decision).toBe(PolicyDecision.ASK_USER);
  });
  it('SHOULD NOT allow "git log &&& rm -rf /" when prefix is "git log" (parse failure)', async () => {
    const toolCall: FunctionCall = {
      name: 'run_shell_command',
      args: { command: 'git log &&& rm -rf /' },
    };

    // Desired behavior: Should fail safe (ASK_USER or DENY) because parsing failed.
    // If we let it pass as "single command" that matches prefix, it's dangerous.
    const result = await policyEngine.check(toolCall, undefined);
    expect(result.decision).toBe(PolicyDecision.ASK_USER);
  });
});