Add basic POC Generation Command to MCP Server t

GEMINI.md

@@ -26,7 +26,7 @@ You are a highly skilled senior security engineer. You are meticulous, an expert
    2. **Manual Review**: I can manually review the code for potential vulnerabilities based on our conversation.
 ```
 *   Explicitly ask the user which they would prefer before proceeding. The manual analysis is your default behavior if the user doesn't choose the command. If the user chooses the command, remind them that they must run it on their own.
-*   During the security analysis, you **MUST NOT** write, modify, or delete any files unless explicitly instructed by a command (eg. `/security:analyze`). Artifacts created during security analysis should be stored in a `.gemini_security/` directory in the user's workspace.
+*   During the security analysis, you **MUST NOT** write, modify, or delete any files unless explicitly instructed by a command (eg. `/security:analyze`). Artifacts created during security analysis should be stored in a `.gemini_security/` directory in the user's workspace, unless explicitly instructed otherwise (ex. `.gemini_security` folder).
 
 ## Skillset: SAST Vulnerability Analysis
 
@@ -192,6 +192,17 @@ For every potential finding, you must perform a quick "So What?" test. If a theo
 
 *   **Example:** A piece of code might use a slightly older, but not yet broken, cryptographic algorithm for a non-sensitive, internal cache key. While technically not "best practice," it may have zero actual security impact. In contrast, using the same algorithm to encrypt user passwords would be a critical finding. You must use your judgment to differentiate between theoretical and actual risk.
 
+### 5. Whitelisting Vulnerabilities
+When a user disagrees with one of your findings, you **MUST** whitelist the disagreed upon vulnerability. 
+
+*   **YOU MUST** Use the MCP Prompt `note-adder` to create a new notation in the `.gemini_security/vuln_whitelist.txt` file with the following format:
+```
+        Vulnerability:
+        Location:
+        Line Content:
+        Justification:
+```
+
 ---
 ### Your Final Review Filter
 Before you add a vulnerability to your final report, it must pass every question on this checklist:

commands/security/analyze.toml

@@ -40,6 +40,8 @@ For EVERY task, you MUST follow this procedure. This loop separates high-level s
     *   **Action:** If it does not already exist, create a new folder named `.gemini_security` in the user's workspace. 
     *   **Action:** Create a new file named `SECURITY_ANALYSIS_TODO.md` in `.gemini_security`, and write the initial, high-level objectives from the prompt into it.
     *   **Action:** Create a new, empty file named `DRAFT_SECURITY_REPORT.md` in `.gemini_security`.
+    *   **Action"** Prep yourself using the following possible notes files under `.gemini_security/`. If they do not exist, skip them.
+        *   `vuln_allowlist.txt`: The allowlist file has vulnerabilities to ignore during your scan. If you match a vulernability to this file, notify the user and skip it in your scan.
 
 2.  **Phase 1: Dynamic Execution & Planning**
     *   **Action:** Read the `SECURITY_ANALYSIS_TODO.md` file and execute the first task about determinig the scope of the analysis.

mcp-server/src/index.ts

@@ -14,6 +14,8 @@ import path from 'path';
 import { getAuditScope } from './filesystem.js';
 import { findLineNumbers } from './security.js';
 
+import { validatePocParams, runPoc } from './poc.js';
+
 const server = new McpServer({
   name: 'gemini-cli-security',
   version: '0.1.0',
@@ -50,6 +52,111 @@ server.tool(
   }
 );
 
+server.tool(
+  'validate_poc_params',
+  'Validates the parameters for the PoC generation.',
+  {
+    vulnerabilityType: z.string().describe('The type of vulnerability.'),
+    sourceCode: z.string().describe('The source code of the vulnerable file.'),
+  },
+  (input) => validatePocParams(input)
+);
+
+server.tool(
+  'run_poc',
+  'Runs the generated PoC code.',
+  {
+    code: z.string().describe('The PoC code to run.'),
+  },
+  (input) => runPoc(input)
+);
+
+server.registerPrompt(
+  'security:note-adder',
+  {
+    title: 'Note Adder',
+    description: 'Creates a new note file or adds a new entry to an existing one, ensuring content consistency.',
+    argsSchema: {
+      notePath: z.string().describe('The path to the note file.'),
+      content: z.string().describe('The content of the note entry to add.'),
+    },
+  },
+  ({ notePath, content }) => ({
+    messages: [
+      {
+        role: 'user',
+        content: {
+          type: 'text',
+          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}'.
+
+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.
+
+2.  **Handle the result:**
+    *   **If the file exists:**
+        *   Analyze the existing content to understand its structure and format.
+        *   **Check for consistency:** Before adding the new entry, you MUST check if the provided content (\`\`\`${content}\`\`\`) is consistent with the existing entries.
+        *   **If it is not consistent:** You MUST ask the user for clarification. Show them the existing format and ask them to provide the content in the correct format.
+        *   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. 
+        *   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.
+
+Your primary goal is to maintain strict consistency with the format of the note file. Do not introduce any formatting changes.`,
+        },
+      },
+    ],
+  }),
+);
+
+server.registerPrompt(
+  'poc',
+  {
+    title: 'PoC Generator',
+    description: 'Generates a Proof-of-Concept (PoC) for a given vulnerability.',
+    argsSchema: {
+      vulnerabilityType: z.string().optional().describe('The type of vulnerability.'),
+      sourceCode: z.string().optional().describe('The source code of the vulnerable file.'),
+    },
+  },
+  ({ vulnerabilityType, sourceCode }) => ({
+    messages: [
+      {
+        role: 'user',
+        content: {
+          type: 'text',
+          text: `You are a security expert. Your task is to generate a Proof-of-Concept (PoC) for a vulnerability.
+
+Input Parameters:
+- Vulnerability Type: ${vulnerabilityType || 'Not provided'}
+- Source Code: ${sourceCode || 'Not provided'}
+
+**Workflow:**
+
+1.  **Validate Parameters:**
+    *   If 'vulnerabilityType' or 'sourceCode' are NOT provided, you MUST ask the user for them.
+    *   If they ARE provided, you MUST use the 'validate_poc_params' tool to validate them.
+
+2.  **Generate PoC:**
+    *   Once parameters are validated, generate a Node.js script that demonstrates the vulnerability under the '.gemini_security/poc/' directory; create the directory if it doesn't exist.
+    *   The script should be self-contained and executable.
+
+3.  **Run PoC:**
+    *   Ask the user for confirmation to run the generated PoC.
+    *   If confirmed, use the 'run_poc' tool with absolute file paths to execute the code.
+    *   Analyze the output to verify if the vulnerability is reproducible.`,
+        },
+      },
+    ],
+  }),
+);
+
 async function startServer() {
   const transport = new StdioServerTransport();
   await server.connect(transport);

mcp-server/src/poc.test.ts

@@ -0,0 +1,88 @@
+import { describe, it, vi, expect } from 'vitest';
+import { validatePocParams, runPoc } from './poc.js';
+
+describe('validatePocParams', () => {
+  it('should return valid message when parameters are provided', async () => {
+    const result = await validatePocParams({
+      vulnerabilityType: 'SQL Injection',
+      sourceCode: 'SELECT * FROM users WHERE id = ' + 1,
+    });
+
+    expect(result.isError).toBeUndefined();
+    expect(result.content[0].text).toBe(
+      JSON.stringify({ message: 'Parameters are valid.' })
+    );
+  });
+
+  it('should return error when vulnerabilityType is missing', async () => {
+    const result = await validatePocParams({
+      vulnerabilityType: '',
+      sourceCode: 'code',
+    });
+
+    expect(result.isError).toBe(true);
+    expect(result.content[0].text).toBe(
+      JSON.stringify({ error: 'Vulnerability type is required.' })
+    );
+  });
+
+  it('should return error when sourceCode is missing', async () => {
+    const result = await validatePocParams({
+      vulnerabilityType: 'type',
+      sourceCode: '',
+    });
+
+    expect(result.isError).toBe(true);
+    expect(result.content[0].text).toBe(
+      JSON.stringify({ error: 'Source code is required.' })
+    );
+  });
+});
+
+describe('runPoc', () => {
+  it('should write file and execute it', async () => {
+    const mockFs = {
+      mkdir: vi.fn(async () => undefined),
+      writeFile: vi.fn(async () => undefined),
+    };
+    const mockPath = {
+      join: (...args: string[]) => args.join('/'),
+    };
+    const mockExecAsync = vi.fn(async () => ({ stdout: 'output', stderr: '' }));
+
+    const result = await runPoc(
+      { code: 'console.log("test")' },
+      { fs: mockFs as any, path: mockPath as any, execAsync: mockExecAsync as any }
+    );
+
+    expect(mockFs.mkdir).toHaveBeenCalledTimes(1);
+    expect(mockFs.writeFile).toHaveBeenCalledTimes(1);
+    expect(mockExecAsync).toHaveBeenCalledTimes(2);
+    expect(result.content[0].text).toBe(
+      JSON.stringify({ stdout: 'output', stderr: '' })
+    );
+  });
+
+  it('should handle execution errors', async () => {
+    const mockFs = {
+      mkdir: vi.fn(async () => undefined),
+      writeFile: vi.fn(async () => undefined),
+    };
+    const mockPath = {
+      join: (...args: string[]) => args.join('/'),
+    };
+    const mockExecAsync = vi.fn(async () => {
+      throw new Error('Execution failed');
+    });
+
+    const result = await runPoc(
+      { code: 'error' },
+      { fs: mockFs as any, path: mockPath as any, execAsync: mockExecAsync as any }
+    );
+
+    expect(result.isError).toBe(true);
+    expect(result.content[0].text).toBe(
+      JSON.stringify({ error: 'Execution failed' })
+    );
+  });
+});

mcp-server/src/poc.ts

@@ -0,0 +1,112 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
+import { promises as fs } from 'fs';
+import path from 'path';
+import { exec } from 'child_process';
+import { promisify } from 'util';
+
+const execAsync = promisify(exec);
+
+export async function validatePocParams(
+  {
+    vulnerabilityType,
+    sourceCode,
+  }: {
+    vulnerabilityType: string;
+    sourceCode: string;
+  }
+): Promise<CallToolResult> {
+  if (!vulnerabilityType || !vulnerabilityType.trim()) {
+    return {
+      content: [
+        {
+          type: 'text',
+          text: JSON.stringify({ error: 'Vulnerability type is required.' }),
+        },
+      ],
+      isError: true,
+    };
+  }
+
+  if (!sourceCode || !sourceCode.trim()) {
+    return {
+      content: [
+        {
+          type: 'text',
+          text: JSON.stringify({ error: 'Source code is required.' }),
+        },
+      ],
+      isError: true,
+    };
+  }
+
+  return {
+    content: [
+      {
+        type: 'text',
+        text: JSON.stringify({ message: 'Parameters are valid.' }),
+      },
+    ],
+  };
+}
+
+export async function runPoc(
+  {
+    code,
+  }: {
+    code: string;
+  },
+  dependencies: { fs: typeof fs; path: typeof path; execAsync: typeof execAsync } = { fs, path, execAsync }
+): Promise<CallToolResult> {
+  try {
+    const CWD = process.cwd();
+    const securityDir = dependencies.path.join(CWD, '.gemini_security');
+
+    // Ensure .gemini_security directory exists
+    try {
+      await dependencies.fs.mkdir(securityDir, { recursive: true });
+    } catch (error) {
+      // Ignore error if directory already exists
+    }
+
+    const pocFilePath = dependencies.path.join(securityDir, 'poc.js');
+    await dependencies.fs.writeFile(pocFilePath, code, 'utf-8');
+
+
+    try {
+      await dependencies.execAsync('npm install', { cwd: securityDir });
+    } 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.
+    }
+    const { stdout, stderr } = await dependencies.execAsync(`node ${pocFilePath}`);
+
+    return {
+      content: [
+        {
+          type: 'text',
+          text: JSON.stringify({ stdout, stderr }),
+        },
+      ],
+    };
+  } catch (error) {
+    let errorMessage = 'An unknown error occurred.';
+    if (error instanceof Error) {
+      errorMessage = error.message;
+    }
+    return {
+      content: [
+        {
+          type: 'text',
+          text: JSON.stringify({ error: errorMessage }),
+        },
+      ],
+      isError: true,
+    };
+  }
+}