diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index aab2c02..6f808f3 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -161,7 +161,7 @@ server.registerPrompt( role: 'user' as const, content: { type: 'text' as const, - text: `You are a security expert. Your task is to generate a Proof-of-Concept (PoC) for a vulnerability. + text: `You are a security expert. Your task is to generate a Proof-of-Concept (PoC) for a vulnerability for Node.js, Python or Go projects. If the project is not for one of these languages, let the user know that you cannot generate a PoC for this project type. Problem Statement: ${problemStatement || 'No problem statement provided, if you need more information to generate a PoC, ask the user.'} Source Code Location: ${sourceCodeLocation || 'No source code location provided, try to derive it from the Problem Statement. If you cannot derive it, ask the user for the source code location.'} @@ -170,7 +170,11 @@ server.registerPrompt( 1. **Generate PoC:** * 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 user's project language, generate a script for Python/Go/Node that demonstrates the vulnerability under the '${SECURITY_DIR_NAME}/${POC_DIR_NAME}/' directory. + * **CRITICAL:** If the PoC script requires external dependencies (e.g. npm packages, PyPI packages) that are not already in the user's project: + * For Node.js: Generate a \`package.json\` in the '${POC_DIR_NAME}' directory. + * For Python: Generate a \`requirements.txt\` in the '${POC_DIR_NAME}' directory. + * For Go: The execution engine will automatically run \`go mod init poc\` and \`go mod tidy\`. * 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. diff --git a/mcp-server/src/poc.test.ts b/mcp-server/src/poc.test.ts index 6a9351d..33c3f79 100644 --- a/mcp-server/src/poc.test.ts +++ b/mcp-server/src/poc.test.ts @@ -16,35 +16,60 @@ describe('runPoc', () => { if (p2) return p1 + '/' + p2; return p1; }, + join: (...paths: string[]) => paths.join('/'), + extname: (p: string) => { + const idx = p.lastIndexOf('.'); + return idx !== -1 ? p.substring(idx) : ''; + }, sep: '/', }; - it('should execute the file at the given path', async () => { - 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: '' }; - }); + it('should execute a Node.js file', async () => { + const mockExecAsync = vi.fn(async () => { return { stdout: '', stderr: '' }; }); + const mockExecFileAsync = vi.fn(async () => { return { stdout: 'output', stderr: '' }; }); const result = await runPoc( { filePath: `${POC_DIR}/test.js` }, - { fs: {} as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any } + { fs: { access: vi.fn().mockRejectedValue(new Error()) } as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any } ); expect(mockExecAsync).toHaveBeenCalledTimes(1); - expect(mockExecAsync).toHaveBeenCalledWith( - 'npm install --registry=https://registry.npmjs.org/', - { cwd: POC_DIR } - ); + expect(mockExecAsync).toHaveBeenCalledWith('npm install --registry=https://registry.npmjs.org/', { cwd: POC_DIR }); 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: '' }) + expect((result.content[0] as any).text).toBe(JSON.stringify({ stdout: 'output', stderr: '' })); + }); + + it('should execute a Python file', async () => { + const mockExecAsync = vi.fn(async () => { return { stdout: '', stderr: '' }; }); + const mockExecFileAsync = vi.fn(async () => { return { stdout: 'output', stderr: '' }; }); + + const result = await runPoc( + { filePath: `${POC_DIR}/test.py` }, + { fs: { access: vi.fn().mockRejectedValue(new Error()) } as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any } + ); + + expect(mockExecAsync).toHaveBeenCalledWith(expect.stringContaining('python3 -m venv')); + expect(mockExecFileAsync).toHaveBeenCalledTimes(1); + expect(mockExecFileAsync).toHaveBeenCalledWith(expect.stringContaining('python'), [`${POC_DIR}/test.py`]); + expect((result.content[0] as any).text).toBe(JSON.stringify({ stdout: 'output', stderr: '' })); + }); + + it('should execute a Go file', async () => { + const mockExecAsync = vi.fn(async () => { return { stdout: '', stderr: '' }; }); + const mockExecFileAsync = vi.fn(async () => { return { stdout: 'output', stderr: '' }; }); + + const result = await runPoc( + { filePath: `${POC_DIR}/test.go` }, + { fs: { access: vi.fn().mockRejectedValue(new Error()) } as any, path: mockPath as any, execAsync: mockExecAsync as any, execFileAsync: mockExecFileAsync as any } ); + + expect(mockExecAsync).toHaveBeenCalledTimes(2); + expect(mockExecAsync).toHaveBeenNthCalledWith(1, 'go mod init poc', { cwd: POC_DIR }); + expect(mockExecAsync).toHaveBeenNthCalledWith(2, 'go mod tidy', { cwd: POC_DIR }); + expect(mockExecFileAsync).toHaveBeenCalledTimes(1); + expect(mockExecFileAsync).toHaveBeenCalledWith('go', ['run', `${POC_DIR}/test.go`]); + expect((result.content[0] as any).text).toBe(JSON.stringify({ stdout: 'output', stderr: '' })); }); it('should handle execution errors', async () => { diff --git a/mcp-server/src/poc.ts b/mcp-server/src/poc.ts index 898c5e2..c48667b 100644 --- a/mcp-server/src/poc.ts +++ b/mcp-server/src/poc.ts @@ -25,7 +25,7 @@ export async function runPoc( try { const pocDir = dependencies.path.dirname(filePath); - // 🛡️ Validate that the filePath is within the safe PoC directory +// 🛡️ 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); @@ -43,13 +43,111 @@ export async function runPoc( }; } + const ext = dependencies.path.extname(filePath).toLowerCase(); + + let installCmd: string | null = null; + let runCmd: string; + let runArgs: string[]; + + if (ext === '.py') { + const venvDir = dependencies.path.join(pocDir, '.venv'); + const isWindows = process.platform === 'win32'; + const pythonBin = isWindows + ? dependencies.path.join(venvDir, 'Scripts', 'python.exe') + : dependencies.path.join(venvDir, 'bin', 'python'); + + try { + await dependencies.fs.access(pythonBin); + } catch { + try { + await dependencies.execAsync(`python3 -m venv "${venvDir}"`); + } catch { + await dependencies.execAsync(`python -m venv "${venvDir}"`); + } + } + + runCmd = pythonBin; + runArgs = [filePath]; + + const projectRoot = process.cwd(); + const checkExists = async (p: string) => + dependencies.fs.access(p).then(() => true).catch(() => false); + + const hasProjectPyproject = await checkExists(dependencies.path.join(projectRoot, 'pyproject.toml')); + const hasProjectRequirements = await checkExists(dependencies.path.join(projectRoot, 'requirements.txt')); + + if (hasProjectPyproject) { + await dependencies.execAsync(`"${pythonBin}" -m pip install -e "${projectRoot}"`).catch(() => { }); + } else if (hasProjectRequirements) { + await dependencies.execAsync(`"${pythonBin}" -m pip install -r "${dependencies.path.join(projectRoot, 'requirements.txt')}"`).catch(() => { }); + } + + const hasPocPyproject = await checkExists(dependencies.path.join(pocDir, 'pyproject.toml')); + const hasPocRequirements = await checkExists(dependencies.path.join(pocDir, 'requirements.txt')); + + if (hasPocPyproject) { + await dependencies.execAsync(`"${pythonBin}" -m pip install .`, { cwd: pocDir }).catch(() => { }); + } + if (hasPocRequirements) { + await dependencies.execAsync(`"${pythonBin}" -m pip install -r requirements.txt`, { cwd: pocDir }).catch(() => { }); + } + } else if (ext === '.go') { + runCmd = 'go'; + runArgs = ['run', filePath]; + + const hasGoMod = await dependencies.fs.access(dependencies.path.join(pocDir, 'go.mod')).then(() => true).catch(() => false); + if (!hasGoMod) { + await dependencies.execAsync('go mod init poc', { cwd: pocDir }).catch(() => { }); + } + + installCmd = 'go mod tidy'; + } else { + runCmd = 'node'; + runArgs = [filePath]; + installCmd = 'npm install --registry=https://registry.npmjs.org/'; + } + + if (installCmd) { + try { + await dependencies.execAsync(installCmd, { cwd: pocDir }); + } catch (error) { + // Ignore errors from install step, as it might fail if no dependency configuration file (e.g., package.json, requirements.txt, go.mod) exists, + // but we still want to attempt running the PoC. + } + } + + let output: { stdout: string; stderr: string }; + 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, - // but we still want to attempt running the PoC. + output = await dependencies.execFileAsync(runCmd, runArgs); + } catch (error: any) { + const errorMessage = error.message || ''; + const errorOutput = (error.stdout || '') + (error.stderr || ''); + + // If we are running a Python script in a venv and it fails due to missing modules, + // try enabling system site packages for the venv and retry. + if (ext === '.py' && (errorMessage.includes('ModuleNotFoundError') || errorOutput.includes('ModuleNotFoundError'))) { + try { + const venvDir = dependencies.path.join(pocDir, '.venv'); + // Update the venv to include system site packages + try { + await dependencies.execAsync(`python3 -m venv --system-site-packages "${venvDir}"`); + } catch { + await dependencies.execAsync(`python -m venv --system-site-packages "${venvDir}"`); + } + + // Retry execution with the updated venv + output = await dependencies.execFileAsync(runCmd, runArgs); + } catch (retryError: any) { + // If retry fails, throw the original error (or the retry error if it's new/different) + throw retryError; + } + } else { + throw error; + } } - const { stdout, stderr } = await dependencies.execFileAsync('node', [filePath]); + + const { stdout, stderr } = output; return { content: [ @@ -61,14 +159,21 @@ export async function runPoc( }; } catch (error) { let errorMessage = 'An unknown error occurred.'; + let stdout = ''; + let stderr = ''; + if (error instanceof Error) { errorMessage = error.message; + // Capture stdout/stderr from the error object if available (execFile throws with these) + stdout = (error as any).stdout || ''; + stderr = (error as any).stderr || ''; } + return { content: [ { type: 'text', - text: JSON.stringify({ error: errorMessage }), + text: JSON.stringify({ error: errorMessage, stdout, stderr }), }, ], isError: true,