diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index 5ac5e17..9f3988f 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,37 @@ 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.'), + sourceCodeLocation: z.string().optional().describe('The location of the source code that contains the vulnerability.'), } as any, }, (args: any) => { - const { vulnerabilityType, sourceCodeLocation } = args; + const { problemStatement, sourceCodeLocation } = 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. - - Input Parameters: - - Vulnerability Type: ${vulnerabilityType || 'Not provided'} - - Source Code Location: ${sourceCodeLocation || 'Not provided'} - - **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.`, + 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 || '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:** + * 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.`, + }, }, - }, - ], - } + ], + }; }, );