Add basic POC Generation Command to MCP Server t #109
Conversation
…ge in security extension
QuinnDACollins
requested review from
QuanZhang-William,
capachino,
evanotero,
heltonduarte,
pedrour and
shrishabh
as code owners
|
🤖 Hi @QuinnDACollins, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This Pull Request introduces a new feature for Proof-of-Concept (PoC) generation and vulnerability whitelisting. The changes include new server tools, prompts, and a poc.ts module with corresponding tests. The overall structure and test coverage are good, and the documentation for whitelisting vulnerabilities is a welcome addition.
🔍 General Feedback
- The prompt for the
note-adderis well-designed, ensuring consistency in note file formats. - The
validatePocParamsfunction provides essential input validation for PoC generation. - Unit tests for
poc.tsare comprehensive for the basic functionality, covering both valid and error scenarios.
| 1. **Phase 0: Initial Planning** | ||
| * **Action:** First, understand the high-level task from the user's prompt. | ||
| * **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. |
There was a problem hiding this comment.
🟢 Low There's a typo in the action description: "Action" should be "Action:".
| * **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:** Prep yourself using the following possible notes files under `.gemini_security/`. If they do not exist, skip them. |
| * **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. |
There was a problem hiding this comment.
🟡 Medium Inconsistency in file naming: the documentation (GEMINI.md) refers to vuln_whitelist.txt, while this TOML file refers to vuln_allowlist.txt. Please standardize on one name for clarity and consistency.
| * **Action"** Prep yourself using the following possible notes files under `.gemini_security/`. If they do not exist, skip them. | |
| * `vuln_whitelist.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. |
There was a problem hiding this comment.
The pull request adds a proof-of-concept generation command to the MCP server. The implementation introduces a critical command injection vulnerability that needs to be addressed.
🔍 General Feedback
- The code is well-structured and easy to understand.
- The addition of tests for the new functionality is a good practice.
| 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.'), |
There was a problem hiding this comment.
is this the actual source code or a source code location?
Shouldn't it be a source code location since we'll likely have that available?
| ); | ||
|
|
||
| server.registerPrompt( | ||
| 'poc', |
There was a problem hiding this comment.
should this be security:poc?
| 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. |
There was a problem hiding this comment.
why is a tool needed to validate the params?
| } | ||
|
|
||
| const pocFilePath = dependencies.path.join(securityDir, 'poc.js'); | ||
| await dependencies.fs.writeFile(pocFilePath, code, 'utf-8'); |
There was a problem hiding this comment.
A bit confused, doesn't the prompt already create the poc.js file?
| // Ensure .gemini_security directory exists | ||
| try { | ||
| await dependencies.fs.mkdir(securityDir, { recursive: true }); | ||
| } catch (error) { |
There was a problem hiding this comment.
I believe mkDir won't error if the dir already exists.
| ``` | ||
| * 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). |
There was a problem hiding this comment.
I think you need to rebase or revert this file? This might be a leftover from an earlier branch/commit.
| 'poc', | ||
| { | ||
| title: 'PoC Generator', | ||
| description: 'Generates a Proof-of-Concept (PoC) for a given vulnerability.', |
There was a problem hiding this comment.
can you maybe add an [Experimental] or something like that to this description.
5 participants
This PR adds a basic poc generation command to the mcp server and the relevant tooling to validate and run the poc.