Restrict PoC command to only run within it's directory, and on a spec…

mcp-server/src/constants.ts

@@ -0,0 +1,13 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import path from 'path';
+
+export const SECURITY_DIR_NAME = '.gemini_security';
+export const POC_DIR_NAME = 'poc';
+
+export const SECURITY_DIR = path.join(process.cwd(), SECURITY_DIR_NAME);
+export const POC_DIR = path.join(SECURITY_DIR, POC_DIR_NAME);

mcp-server/src/index.ts

@@ -14,6 +14,7 @@ import path from 'path';
 import { getAuditScope } from './filesystem.js';
 import { findLineNumbers } from './security.js';
 import { parseMarkdownToDict } from './parser.js';
+import { SECURITY_DIR_NAME, POC_DIR_NAME } from './constants.js';
 
 import { runPoc } from './poc.js';
 
@@ -22,57 +23,16 @@ const server = new McpServer({
   version: '0.1.0',
 });
 
-server.tool(
-  'find_line_numbers',
-  'Finds the line numbers of a code snippet in a file.',
-  {
-    filePath: z
-      .string()
-      .describe('The path to the file to with the security vulnerability.'),
-    snippet: z
-      .string()
-      .describe('The code snippet to search for inside the file.'),
-  } as any,
-  (input: { filePath: string; snippet: string }) => findLineNumbers(input, { fs, path })
-);
-
-server.tool(
-  'get_audit_scope',
-  'Gets the git diff of the current changes. Can optionally compare two specific branches.',
-  {
-    base: z.string().optional().describe('The base branch or commit hash (e.g., "main").'),
-    head: z.string().optional().describe('The head branch or commit hash (e.g., "feature-branch").'),
-  } as any,
-  ((args: { base?: string; head?: string }) => {
-    const diff = getAuditScope(args.base, args.head);
-    return {
-      content: [
-        {
-          type: 'text',
-          text: diff,
-        },
-      ],
-    };
-  }) as any
-);
-
-server.tool(
-  'run_poc',
-  'Runs the generated PoC code.',
-  {
-    filePath: z.string().describe('The absolute path to the PoC file to run.'),
-  } as any,
-  (input: { filePath: string }) => runPoc(input)
-);
+// ... (omitted lines)
 
 server.tool(
   'convert_report_to_json',
-  'Converts the Markdown security report into a JSON file named security_report.json in the .gemini_security folder.',
+  `Converts the Markdown security report into a JSON file named security_report.json in the ${SECURITY_DIR_NAME} folder.`,
   {} as any,
   (async () => {
     try {
-      const reportPath = path.join(process.cwd(), '.gemini_security/DRAFT_SECURITY_REPORT.md');
-      const outputPath = path.join(process.cwd(), '.gemini_security/security_report.json');
+      const reportPath = path.join(process.cwd(), `${SECURITY_DIR_NAME}/DRAFT_SECURITY_REPORT.md`);
+      const outputPath = path.join(process.cwd(), `${SECURITY_DIR_NAME}/security_report.json`);
 
       const content = await fs.readFile(reportPath, 'utf-8');
       const results = parseMarkdownToDict(content);
@@ -113,13 +73,13 @@ server.registerPrompt(
         role: 'user' as const,
         content: {
           type: 'text' as const,
-          text: `You are a helpful assistant that helps users maintain notes. Your task is to add a new entry to the notes file at '.gemini_security/${notePath}'.
+          text: `You are a helpful assistant that helps users maintain notes. Your task is to add a new entry to the notes file at '${SECURITY_DIR_NAME}/${notePath}'.
 
         You MUST use the 'ReadFile' and 'WriteFile' tools.
 
         **Workflow:**
 
-        1.  **Read the file:** First, you MUST attempt to read the file at '.gemini_security/${notePath}' using the 'ReadFile' tool.
+        1.  **Read the file:** First, you MUST attempt to read the file at '${SECURITY_DIR_NAME}/${notePath}' using the 'ReadFile' tool.
 
         2.  **Handle the result:**
             *   **If the file exists:**
@@ -129,7 +89,7 @@ server.registerPrompt(
                 *   Once you have a consistent entry, append it to the content, ensuring it perfectly matches the existing format.
                 *   Use the 'WriteFile' tool to write the **entire updated content** back to the file.
             *   **If the file does NOT exist (ReadFile returns an error):**
-                *   First, if the '.gemini_security' directory doesn't exist, create it. 
+                *   First, if the '${SECURITY_DIR_NAME}' directory doesn't exist, create it. 
                 *   This is a new note. You MUST ask the user to define a template for this note.
                 *   Once the user provides a template, construct the initial file content. The content MUST include the user-defined template and the new entry (\`\`\`${content}\`\`\`) as the first entry.
                 *   Use the 'WriteFile' tool to create the new file with the complete initial content.
@@ -168,8 +128,13 @@ server.registerPrompt(
           **Workflow:**
 
           1.  **Generate PoC:**
-              *   Create a 'poc' directory in '.gemini_security' if it doesn't exist.
-              *   Generate a Node.js script that demonstrates the vulnerability under the '.gemini_security/poc/' directory.
+              *   Create a '${POC_DIR_NAME}' directory in '${SECURITY_DIR_NAME}' if it doesn't exist.
+              *   Generate a Node.js script that demonstrates the vulnerability under the '${SECURITY_DIR_NAME}/${POC_DIR_NAME}/' directory.
+              *   Based on the vulnerability type certain criteria must be met in our script, otherwise generate the PoC to the best of your ability:
+                  *   If the vulnerability is a Path Traversal Vulnerability:
+                      *   **YOU MUST** Use the 'write_file' tool to create a temporary file '../gcli_secext_temp.txt' directly outside of the project directory. 
+                      *   The script should then try to read the file using the vulnerability in the user's code. 
+                      *   **YOU MUST** Use the 'write_file' tool to delete the '../gcli_secext_temp.txt' file after the verification step, regardless of whether the read was successful or not.
               *   The script should import the user's vulnerable file(s), and demonstrate the vulnerability in their code.
 
           2.  **Run PoC:**

mcp-server/src/poc.test.ts

@@ -6,55 +6,78 @@
 
 import { describe, it, vi, expect } from 'vitest';
 import { runPoc } from './poc.js';
+import { POC_DIR } from './constants.js';
 
 describe('runPoc', () => {
+  const mockPath = {
+    dirname: (p: string) => p.substring(0, p.lastIndexOf('/')),
+    resolve: (p1: string, p2?: string) => {
+      if (p2 && p2.startsWith('/')) return p2;
+      if (p2) return p1 + '/' + p2;
+      return p1;
+    },
+    sep: '/',
+  };
+
   it('should execute the file at the given path', async () => {
-    const mockPath = {
-      dirname: (p: string) => p.substring(0, p.lastIndexOf('/')),
-    };
     const mockExecAsync = vi.fn(async (cmd: string) => {
       if (cmd.startsWith('npm install')) {
         return { stdout: '', stderr: '' };
       }
       return { stdout: 'output', stderr: '' };
     });
+    const mockExecFileAsync = vi.fn(async (file: string, args?: string[]) => {
+      return { stdout: 'output', stderr: '' };
+    });
 
     const result = await runPoc(
-      { filePath: '/tmp/test.js' },
-      { fs: {} as any, path: mockPath as any, execAsync: mockExecAsync as any }
+      { filePath: `${POC_DIR}/test.js` },
+      { fs: {} as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any }
     );
 
-    expect(mockExecAsync).toHaveBeenCalledTimes(2);
-    expect(mockExecAsync).toHaveBeenNthCalledWith(
-      1,
+    expect(mockExecAsync).toHaveBeenCalledTimes(1);
+    expect(mockExecAsync).toHaveBeenCalledWith(
       'npm install --registry=https://registry.npmjs.org/',
-      { cwd: '/tmp' }
+      { cwd: POC_DIR }
     );
-    expect(mockExecAsync).toHaveBeenNthCalledWith(2, 'node /tmp/test.js');
+    expect(mockExecFileAsync).toHaveBeenCalledTimes(1);
+    expect(mockExecFileAsync).toHaveBeenCalledWith('node', [`${POC_DIR}/test.js`]);
     expect((result.content[0] as any).text).toBe(
       JSON.stringify({ stdout: 'output', stderr: '' })
     );
   });
 
   it('should handle execution errors', async () => {
-    const mockPath = {
-      dirname: (p: string) => p.substring(0, p.lastIndexOf('/')),
-    };
     const mockExecAsync = vi.fn(async (cmd: string) => {
-      if (cmd.startsWith('node')) {
-        throw new Error('Execution failed');
-      }
       return { stdout: '', stderr: '' };
     });
+    const mockExecFileAsync = vi.fn(async (file: string, args?: string[]) => {
+      throw new Error('Execution failed');
+    });
 
     const result = await runPoc(
-      { filePath: '/tmp/error.js' },
-      { fs: {} as any, path: mockPath as any, execAsync: mockExecAsync as any }
+      { filePath: `${POC_DIR}/error.js` },
+      { fs: {} as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any }
     );
 
     expect(result.isError).toBe(true);
     expect((result.content[0] as any).text).toBe(
       JSON.stringify({ error: 'Execution failed' })
     );
   });
+
+  it('should fail when accessing file outside of allowed directory', async () => {
+    const mockExecAsync = vi.fn();
+    const mockExecFileAsync = vi.fn();
+
+    const result = await runPoc(
+      { filePath: '/tmp/malicious.js' },
+      { fs: {} as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any }
+    );
+
+    expect(result.isError).toBe(true);
+    expect((result.content[0] as any).text).toContain('Security Error: PoC execution is restricted');
+    expect(mockExecAsync).not.toHaveBeenCalled();
+    expect(mockExecFileAsync).not.toHaveBeenCalled();
+  });
 });

mcp-server/src/poc.ts

@@ -7,29 +7,49 @@
 import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import { promises as fs } from 'fs';
 import path from 'path';
-import { exec } from 'child_process';
+import { exec, execFile } from 'child_process';
 import { promisify } from 'util';
+import { POC_DIR } from './constants.js';
 
 const execAsync = promisify(exec);
+const execFileAsync = promisify(execFile);
 
 export async function runPoc(
   {
     filePath,
   }: {
       filePath: string;
   },
-  dependencies: { fs: typeof fs; path: typeof path; execAsync: typeof execAsync } = { fs, path, execAsync }
+  dependencies: { fs: typeof fs; path: typeof path; execAsync: typeof execAsync; execFileAsync: typeof execFileAsync } = { fs, path, execAsync, execFileAsync }
 ): Promise<CallToolResult> {
   try {
     const pocDir = dependencies.path.dirname(filePath);
 
+    // 🛡️ Validate that the filePath is within the safe PoC directory
+    const resolvedFilePath = dependencies.path.resolve(filePath);
+    const safePocDir = dependencies.path.resolve(process.cwd(), POC_DIR);
+
+    if (!resolvedFilePath.startsWith(safePocDir + dependencies.path.sep)) {
+      return {
+        content: [
+          {
+            type: 'text',
+            text: JSON.stringify({
+              error: `Security Error: PoC execution is restricted to files within '${safePocDir}'. Attempted to access '${resolvedFilePath}'.`,
+            }),
+          },
+        ],
+        isError: true,
+      };
+    }
+
     try {
       await dependencies.execAsync('npm install --registry=https://registry.npmjs.org/', { cwd: pocDir });
     } catch (error) {
-      // Ignore errors from npm install, as it might fail if no package.json exists,
+      // 📦 Ignore errors from npm install, as it might fail if no package.json exists,
       // but we still want to attempt running the PoC.
     }
-    const { stdout, stderr } = await dependencies.execAsync(`node ${filePath}`);
+    const { stdout, stderr } = await dependencies.execFileAsync('node', [filePath]);
 
     return {
       content: [