From a0449d3baddc9833bdca68af91eca27446c83c2c Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Tue, 16 Dec 2025 11:32:36 -0800 Subject: [PATCH 1/2] feature: Use problem statements in the PoC function to allow for more flexible usage --- mcp-server/src/index.ts | 91 +++++++++++++++++++++-------------------- 1 file changed, 47 insertions(+), 44 deletions(-) diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index 5ac5e17..6c9631f 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -81,26 +81,26 @@ server.registerPrompt( 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}'. -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.`, + 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.`, }, }, ], @@ -114,39 +114,42 @@ server.registerPrompt( title: 'PoC Generator', description: '[Experimental] Generates a Proof-of-Concept (PoC) for a given vulnerability.', argsSchema: { - vulnerabilityType: z.string().optional().describe('The type of vulnerability.'), - sourceCodeLocation: z.string().optional().describe('The location of the source code of the vulnerable file.'), + problemStatement: z.string().optional().describe('A description of the security problem or vulnerability.'), } as any, }, (args: any) => { - const { vulnerabilityType, sourceCodeLocation } = args; + const { problemStatement } = args; return { - messages: [ - { - 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. - Use the given parameters to generate the PoC, if they don't exist, ask the user to provide them. + messages: [ + { + 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. + + Problem Statement: ${problemStatement || 'Not provided'} - Input Parameters: - - Vulnerability Type: ${vulnerabilityType || 'Not provided'} - - Source Code Location: ${sourceCodeLocation || 'Not provided'} + Using the provided problem statement, identify the following parameters: + - Source Code Location - **Workflow:** + If the problem statement does not provide enough information to identify the parameters, **ASK THE USER** for the missing information. - 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. - * The script should import the user's vulnerable file(s), and demonstrate the vulnerability in their code. + If you have all the necessary information, proceed with the following workflow: - 2. **Run PoC:** - * Use the 'run_poc' tool with absolute file paths to execute the code. - * Analyze the output to verify if the vulnerability is reproducible.`, + **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. + * The script should import the user's vulnerable file(s), and demonstrate the vulnerability in their code. + + 2. **Run PoC:** + * Use the 'run_poc' tool with absolute file paths to execute the code. + * Analyze the output to verify if the vulnerability is reproducible.`, + }, }, - }, - ], - } + ], + }; }, ); From 07a1e07e215b592f261e6d58b6b98a9e8bf86c31 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Wed, 17 Dec 2025 12:31:06 -0800 Subject: [PATCH 2/2] fix: Add source code location as parmater to our PoC command (it increases success rate and decreases token usage) --- mcp-server/src/index.ts | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index 6c9631f..9f3988f 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -115,10 +115,11 @@ server.registerPrompt( description: '[Experimental] Generates a Proof-of-Concept (PoC) for a given vulnerability.', argsSchema: { problemStatement: z.string().optional().describe('A description of the security problem or vulnerability.'), + sourceCodeLocation: z.string().optional().describe('The location of the source code that contains the vulnerability.'), } as any, }, (args: any) => { - const { problemStatement } = args; + const { problemStatement, sourceCodeLocation } = args; return { messages: [ { @@ -127,15 +128,9 @@ server.registerPrompt( type: 'text' as const, text: `You are a security expert. Your task is to generate a Proof-of-Concept (PoC) for a vulnerability. - Problem Statement: ${problemStatement || 'Not provided'} - - Using the provided problem statement, identify the following parameters: - - Source Code Location - - If the problem statement does not provide enough information to identify the parameters, **ASK THE USER** for the missing information. - - If you have all the necessary information, proceed with the following workflow: - + 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.'} + **Workflow:** 1. **Generate PoC:**