diff --git a/mcp-server/src/constants.ts b/mcp-server/src/constants.ts new file mode 100644 index 0000000..4ed3305 --- /dev/null +++ b/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); diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index 86544ac..2a0edc8 100644 --- a/mcp-server/src/index.ts +++ b/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:** diff --git a/mcp-server/src/poc.test.ts b/mcp-server/src/poc.test.ts index 70decb4..6a9351d 100644 --- a/mcp-server/src/poc.test.ts +++ b/mcp-server/src/poc.test.ts @@ -6,50 +6,58 @@ 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); @@ -57,4 +65,19 @@ describe('runPoc', () => { 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(); + }); }); diff --git a/mcp-server/src/poc.ts b/mcp-server/src/poc.ts index d75d740..898c5e2 100644 --- a/mcp-server/src/poc.ts +++ b/mcp-server/src/poc.ts @@ -7,10 +7,12 @@ 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( { @@ -18,18 +20,36 @@ export async function runPoc( }: { 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 { 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: [