From 78bd06a806d16983ccb2f8bf9f855a87711dd83c Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Tue, 11 Nov 2025 12:16:00 -0800 Subject: [PATCH 01/11] Add registered prompt and gemini md updates for enabling notation usage in security extension --- GEMINI.md | 12 +++++++++- commands/security/analyze.toml | 2 ++ mcp-server/src/index.ts | 43 ++++++++++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/GEMINI.md b/GEMINI.md index 931e80c..9e1bc84 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -193,7 +193,17 @@ For every potential finding, you must perform a quick "So What?" test. If a theo * **YOU MUST** Only perform security analysis when the user explicitly requests for help with code security or vulnerabilities. Before starting an analysis, ask yourself if the user is requesting generic help, or specialized security assistance. ---- +### 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 `security_notes/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: diff --git a/commands/security/analyze.toml b/commands/security/analyze.toml index 8fccfb6..6bcecf6 100644 --- a/commands/security/analyze.toml +++ b/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 notes files under `security_notes/` + * `vuln_whitelist.txt`: The whitelist 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. diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index aee8db1..a7b287c 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -50,6 +50,49 @@ server.tool( } ); +server.registerPrompt( + '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 note file at ${notePath}. + +You MUST use the 'ReadFile' and 'WriteFile' tools. + +**Workflow:** + +1. **Read the file:** First, you MUST attempt to read the file at ${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):** + * 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.`, + }, + }, + ], + }), +); + async function startServer() { const transport = new StdioServerTransport(); await server.connect(transport); From 32ad411db1f7bb1a6fe3d0a756f982dafbdae6a0 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Wed, 12 Nov 2025 09:35:26 -0800 Subject: [PATCH 02/11] fix: folder location wording in gemini .md --- GEMINI.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/GEMINI.md b/GEMINI.md index 9e1bc84..b8da419 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -25,7 +25,9 @@ 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. `security_notes` folder). + +--- ## Skillset: SAST Vulnerability Analysis @@ -196,7 +198,7 @@ For every potential finding, you must perform a quick "So What?" test. If a theo ### 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 `security_notes/vuln_whitelist.txt` file with the following format: +* **YOU MUST** Use the MCP Prompt `note-adder` to create a new note folder `security_notes`: `security_notes/vuln_whitelist.txt` with the following format: ``` Vulnerability: Location: From 8cbfd3c883cb9a8e3fb80df1a684326805de81f2 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Wed, 12 Nov 2025 11:43:15 -0800 Subject: [PATCH 03/11] fix: merge into main --- GEMINI.md | 16 +++++++++++++ commands/security/analyze.toml | 2 ++ mcp-server/src/index.ts | 43 ++++++++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+) diff --git a/GEMINI.md b/GEMINI.md index e9c36b2..807e033 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -192,7 +192,23 @@ 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. +<<<<<<< HEAD --- +======= +* **YOU MUST** Only perform security analysis when the user explicitly requests for help with code security or vulnerabilities. Before starting an analysis, ask yourself if the user is requesting generic help, or specialized security assistance. + +### 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 `security_notes/vuln_whitelist.txt` file with the following format: +``` + Vulnerability: + Location: + Line Content: + Justification: +``` + +>>>>>>> 78bd06a (Add registered prompt and gemini md updates for enabling notation usage in security extension) ### Your Final Review Filter Before you add a vulnerability to your final report, it must pass every question on this checklist: diff --git a/commands/security/analyze.toml b/commands/security/analyze.toml index 7c3e07f..961be8a 100644 --- a/commands/security/analyze.toml +++ b/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 notes files under `security_notes/` + * `vuln_whitelist.txt`: The whitelist 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. diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index aee8db1..a7b287c 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -50,6 +50,49 @@ server.tool( } ); +server.registerPrompt( + '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 note file at ${notePath}. + +You MUST use the 'ReadFile' and 'WriteFile' tools. + +**Workflow:** + +1. **Read the file:** First, you MUST attempt to read the file at ${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):** + * 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.`, + }, + }, + ], + }), +); + async function startServer() { const transport = new StdioServerTransport(); await server.connect(transport); From da3ef9906ad11efd20af574fe12967e36064dffa Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Wed, 12 Nov 2025 09:35:26 -0800 Subject: [PATCH 04/11] fix: folder location wording in gemini .md --- GEMINI.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/GEMINI.md b/GEMINI.md index 807e033..57b81e3 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -26,7 +26,9 @@ 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. `security_notes` folder). + +--- ## Skillset: SAST Vulnerability Analysis @@ -200,7 +202,7 @@ For every potential finding, you must perform a quick "So What?" test. If a theo ### 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 `security_notes/vuln_whitelist.txt` file with the following format: +* **YOU MUST** Use the MCP Prompt `note-adder` to create a new note folder `security_notes`: `security_notes/vuln_whitelist.txt` with the following format: ``` Vulnerability: Location: From 1c8779066dc8c43c01d5df4990efc0e6f4ab9b83 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Wed, 12 Nov 2025 11:54:12 -0800 Subject: [PATCH 05/11] fix: remove merge remnants --- GEMINI.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/GEMINI.md b/GEMINI.md index 2a1d22a..daa88cb 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -28,8 +28,6 @@ You are a highly skilled senior security engineer. You are meticulous, an expert * 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, unless explicitly instructed otherwise (ex. `security_notes` folder). ---- - ## Skillset: SAST Vulnerability Analysis This is your internal knowledge base of vulnerabilities. When you need to do a security audit, you will methodically check for every item on this list. @@ -205,6 +203,7 @@ When a user disagrees with one of your findings, you **MUST** whitelist the disa Justification: ``` +--- ### Your Final Review Filter Before you add a vulnerability to your final report, it must pass every question on this checklist: From 0ea0b48f9d95dd2a9af977928824dac7141a46e8 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Thu, 13 Nov 2025 14:28:22 -0800 Subject: [PATCH 06/11] fix: make prompt less error prone by enforcing directory --- commands/security/analyze.toml | 2 +- mcp-server/src/index.ts | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/commands/security/analyze.toml b/commands/security/analyze.toml index 961be8a..1e1b499 100644 --- a/commands/security/analyze.toml +++ b/commands/security/analyze.toml @@ -41,7 +41,7 @@ For EVERY task, you MUST follow this procedure. This loop separates high-level s * **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 notes files under `security_notes/` - * `vuln_whitelist.txt`: The whitelist 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. + * `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. diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index a7b287c..949fddd 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -51,7 +51,7 @@ server.tool( ); server.registerPrompt( - 'note-adder', + 'security:note-adder', { title: 'Note Adder', description: 'Creates a new note file or adds a new entry to an existing one, ensuring content consistency.', @@ -66,13 +66,13 @@ server.registerPrompt( 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 note file at ${notePath}. + text: `You are a helpful assistant that helps users maintain notes. Your task is to add a new entry to the notes file at 'security_notes/${notePath}'. You MUST use the 'ReadFile' and 'WriteFile' tools. **Workflow:** -1. **Read the file:** First, you MUST attempt to read the file at ${notePath} using the 'ReadFile' tool. +1. **Read the file:** First, you MUST attempt to read the file at 'security_notes/${notePath}' using the 'ReadFile' tool. 2. **Handle the result:** * **If the file exists:** @@ -82,6 +82,7 @@ You MUST use the 'ReadFile' and 'WriteFile' tools. * 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 'security_notes' 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. From bac4ab6ffaa00d07adb60b8eae98ce8e7d51f43d Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Mon, 17 Nov 2025 09:33:50 -0800 Subject: [PATCH 07/11] fix: move whitelist directory to .gemini_security --- commands/security/analyze.toml | 2 +- mcp-server/src/index.ts | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/commands/security/analyze.toml b/commands/security/analyze.toml index 1e1b499..91b1169 100644 --- a/commands/security/analyze.toml +++ b/commands/security/analyze.toml @@ -40,7 +40,7 @@ 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 notes files under `security_notes/` + * **Action"** Prep yourself using the following notes files under `.gemini_security/` * `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** diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index 949fddd..23bc8a7 100644 --- a/mcp-server/src/index.ts +++ b/mcp-server/src/index.ts @@ -66,13 +66,13 @@ server.registerPrompt( 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 'security_notes/${notePath}'. + 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 'security_notes/${notePath}' using the 'ReadFile' tool. +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:** @@ -82,7 +82,7 @@ You MUST use the 'ReadFile' and 'WriteFile' tools. * 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 'security_notes' directory doesn't exist, create it. + * 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. From 1723ce8db545368adb39629f6c16d33780797770 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Mon, 17 Nov 2025 13:48:35 -0800 Subject: [PATCH 08/11] fix: remove mentions of unused security notes folder from gemini md --- GEMINI.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/GEMINI.md b/GEMINI.md index daa88cb..697a904 100644 --- a/GEMINI.md +++ b/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, unless explicitly instructed otherwise (ex. `security_notes` folder). +* 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 @@ -195,7 +195,7 @@ For every potential finding, you must perform a quick "So What?" test. If a theo ### 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 `security_notes/vuln_whitelist.txt` file with the following format: +* **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: From e0f60ea96da86bf12f272c7c7b3f5c75b1bec113 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Mon, 17 Nov 2025 14:02:43 -0800 Subject: [PATCH 09/11] fix: add language that suggests to skip if note doesnt exist --- commands/security/analyze.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/commands/security/analyze.toml b/commands/security/analyze.toml index 91b1169..c629d65 100644 --- a/commands/security/analyze.toml +++ b/commands/security/analyze.toml @@ -40,7 +40,7 @@ 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 notes files under `.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** From 2f533fdb65368aa64219bd772d0228f73b544c36 Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Fri, 21 Nov 2025 16:28:56 -0800 Subject: [PATCH 10/11] feature: Add basic poc command functionality to the MCP server --- mcp-server/src/index.ts | 63 +++++++++++++++++++++ mcp-server/src/poc.test.ts | 88 +++++++++++++++++++++++++++++ mcp-server/src/poc.ts | 112 +++++++++++++++++++++++++++++++++++++ 3 files changed, 263 insertions(+) create mode 100644 mcp-server/src/poc.test.ts create mode 100644 mcp-server/src/poc.ts diff --git a/mcp-server/src/index.ts b/mcp-server/src/index.ts index 23bc8a7..9bbd631 100644 --- a/mcp-server/src/index.ts +++ b/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,25 @@ 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', { @@ -94,6 +115,48 @@ Your primary goal is to maintain strict consistency with the format of the note }), ); +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); diff --git a/mcp-server/src/poc.test.ts b/mcp-server/src/poc.test.ts new file mode 100644 index 0000000..8ba9cf1 --- /dev/null +++ b/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' }) + ); + }); +}); diff --git a/mcp-server/src/poc.ts b/mcp-server/src/poc.ts new file mode 100644 index 0000000..bd6b888 --- /dev/null +++ b/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 { + 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 { + 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, + }; + } +} From 7e5ea1888ac3cd1ec355f688f9ec3a248ed75ccb Mon Sep 17 00:00:00 2001 From: QuinnDACollins Date: Tue, 2 Dec 2025 11:09:52 -0800 Subject: [PATCH 11/11] fix: use isolated-vm library to isolate generated code --- mcp-server/package-lock.json | 395 +++++++++++++++++++++++++++++++++++ mcp-server/package.json | 1 + mcp-server/src/poc.test.ts | 36 +--- mcp-server/src/poc.ts | 44 ++-- 4 files changed, 424 insertions(+), 52 deletions(-) diff --git a/mcp-server/package-lock.json b/mcp-server/package-lock.json index 73d7315..99965b8 100644 --- a/mcp-server/package-lock.json +++ b/mcp-server/package-lock.json @@ -7,6 +7,7 @@ "name": "gemini-cli-security-mcp-server", "dependencies": { "@modelcontextprotocol/sdk": "^1.18.0", + "isolated-vm": "^6.0.2", "zod": "^3.24.2" }, "devDependencies": { @@ -984,6 +985,37 @@ "node": ">=12" } }, + "node_modules/base64-js": { + "version": "1.5.1", + "resolved": "https://registry.npmjs.org/base64-js/-/base64-js-1.5.1.tgz", + "integrity": "sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "MIT" + }, + "node_modules/bl": { + "version": "4.1.0", + "resolved": "https://registry.npmjs.org/bl/-/bl-4.1.0.tgz", + "integrity": "sha512-1W07cM9gS6DcLperZfFSj+bWLtaPGSOHWhPiGzXmvVJbRLdG82sH/Kn8EtW1VqWVA54AKf2h5k5BbnIbwF3h6w==", + "license": "MIT", + "dependencies": { + "buffer": "^5.5.0", + "inherits": "^2.0.4", + "readable-stream": "^3.4.0" + } + }, "node_modules/body-parser": { "version": "2.2.0", "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-2.2.0.tgz", @@ -1004,6 +1036,30 @@ "node": ">=18" } }, + "node_modules/buffer": { + "version": "5.7.1", + "resolved": "https://registry.npmjs.org/buffer/-/buffer-5.7.1.tgz", + "integrity": "sha512-EHcyIPBQ4BSGlvjB16k5KgAJ27CIsHY/2JBmCRReo48y9rQ3MaUzWX3KVlBa4U7MyX02HdVj0K7C3WaB3ju7FQ==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "MIT", + "dependencies": { + "base64-js": "^1.3.1", + "ieee754": "^1.1.13" + } + }, "node_modules/bytes": { "version": "3.1.2", "resolved": "https://registry.npmjs.org/bytes/-/bytes-3.1.2.tgz", @@ -1079,6 +1135,12 @@ "node": ">= 16" } }, + "node_modules/chownr": { + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/chownr/-/chownr-1.1.4.tgz", + "integrity": "sha512-jJ0bqzaylmJtVnNgzTeSOs8DPavpbYgEr/b0YL8/2GO3xJEhInFmhKMUnEJQjZumK7KXGFhUy89PrsJWlakBVg==", + "license": "ISC" + }, "node_modules/content-disposition": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-1.0.0.tgz", @@ -1162,6 +1224,21 @@ } } }, + "node_modules/decompress-response": { + "version": "6.0.0", + "resolved": "https://registry.npmjs.org/decompress-response/-/decompress-response-6.0.0.tgz", + "integrity": "sha512-aW35yZM6Bb/4oJlZncMH2LCoZtJXTRxES17vE3hoRiowU2kWHaJKFkSBDnDR+cm9J+9QhXmREyIfv0pji9ejCQ==", + "license": "MIT", + "dependencies": { + "mimic-response": "^3.1.0" + }, + "engines": { + "node": ">=10" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, "node_modules/deep-eql": { "version": "5.0.2", "resolved": "https://registry.npmjs.org/deep-eql/-/deep-eql-5.0.2.tgz", @@ -1172,6 +1249,15 @@ "node": ">=6" } }, + "node_modules/deep-extend": { + "version": "0.6.0", + "resolved": "https://registry.npmjs.org/deep-extend/-/deep-extend-0.6.0.tgz", + "integrity": "sha512-LOHxIOaPYdHlJRtCQfDIVZtfw/ufM8+rVj649RIHzcm/vGwQRXFt6OPqIFWsm2XEMrNIEtWR64sY1LEKD2vAOA==", + "license": "MIT", + "engines": { + "node": ">=4.0.0" + } + }, "node_modules/depd": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", @@ -1181,6 +1267,15 @@ "node": ">= 0.8" } }, + "node_modules/detect-libc": { + "version": "2.1.2", + "resolved": "https://registry.npmjs.org/detect-libc/-/detect-libc-2.1.2.tgz", + "integrity": "sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==", + "license": "Apache-2.0", + "engines": { + "node": ">=8" + } + }, "node_modules/dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", @@ -1210,6 +1305,15 @@ "node": ">= 0.8" } }, + "node_modules/end-of-stream": { + "version": "1.4.5", + "resolved": "https://registry.npmjs.org/end-of-stream/-/end-of-stream-1.4.5.tgz", + "integrity": "sha512-ooEGc6HP26xXq/N+GCGOT0JKCLDGrq2bQUZrQ7gyrJiZANJ/8YDTxTpQBXGMn+WbIQXNVpyWymm7KYVICQnyOg==", + "license": "MIT", + "dependencies": { + "once": "^1.4.0" + } + }, "node_modules/es-define-property": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/es-define-property/-/es-define-property-1.0.1.tgz", @@ -1335,6 +1439,15 @@ "node": ">=18.0.0" } }, + "node_modules/expand-template": { + "version": "2.0.3", + "resolved": "https://registry.npmjs.org/expand-template/-/expand-template-2.0.3.tgz", + "integrity": "sha512-XYfuKMvj4O35f/pOXLObndIRvyQ+/+6AhODh+OKWj9S9498pHHn/IMszH+gt0fBCRWMNfk1ZSp5x3AifmnI2vg==", + "license": "(MIT OR WTFPL)", + "engines": { + "node": ">=6" + } + }, "node_modules/expect-type": { "version": "1.2.2", "resolved": "https://registry.npmjs.org/expect-type/-/expect-type-1.2.2.tgz", @@ -1467,6 +1580,12 @@ "node": ">= 0.8" } }, + "node_modules/fs-constants": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/fs-constants/-/fs-constants-1.0.0.tgz", + "integrity": "sha512-y6OAwoSIf7FyjMIv94u+b5rdheZEjzR63GTyZJm5qh4Bi+2YgwLCcI/fPFZkL5PSixOt6ZNKm+w+Hfp/Bciwow==", + "license": "MIT" + }, "node_modules/fsevents": { "version": "2.3.3", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", @@ -1528,6 +1647,12 @@ "node": ">= 0.4" } }, + "node_modules/github-from-package": { + "version": "0.0.0", + "resolved": "https://registry.npmjs.org/github-from-package/-/github-from-package-0.0.0.tgz", + "integrity": "sha512-SyHy3T1v2NUXn29OsWdxmK6RwHD+vkj3v8en8AOBZ1wBQ/hCAQ5bAQTD02kW4W9tUp/3Qh6J8r9EvntiyCmOOw==", + "license": "MIT" + }, "node_modules/gopd": { "version": "1.2.0", "resolved": "https://registry.npmjs.org/gopd/-/gopd-1.2.0.tgz", @@ -1601,12 +1726,38 @@ "node": ">=0.10.0" } }, + "node_modules/ieee754": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/ieee754/-/ieee754-1.2.1.tgz", + "integrity": "sha512-dcyqhDvX1C46lXZcVqCpK+FtMRQVdIMN6/Df5js2zouUsqG7I6sFxitIC+7KYK29KdXOLHdu9zL4sFnoVQnqaA==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "BSD-3-Clause" + }, "node_modules/inherits": { "version": "2.0.4", "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", "license": "ISC" }, + "node_modules/ini": { + "version": "1.3.8", + "resolved": "https://registry.npmjs.org/ini/-/ini-1.3.8.tgz", + "integrity": "sha512-JV/yugV2uzW5iMRSiZAyDtQd+nxtUnjeLt0acNdw98kKLrvuRVyB80tsREOE7yvGVgalhZ6RNXCmEHkUKBKxew==", + "license": "ISC" + }, "node_modules/ipaddr.js": { "version": "1.9.1", "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz", @@ -1628,6 +1779,19 @@ "integrity": "sha512-RHxMLp9lnKHGHRng9QFhRCMbYAcVpn69smSGcq3f36xjgVVWThj4qqLbTLlq7Ssj8B+fIQ1EuCEGI2lKsyQeIw==", "license": "ISC" }, + "node_modules/isolated-vm": { + "version": "6.0.2", + "resolved": "https://registry.npmjs.org/isolated-vm/-/isolated-vm-6.0.2.tgz", + "integrity": "sha512-Qw6AJuagG/VJuh2AIcSWmQPsAArti/L+lKhjXU+lyhYkbt3J57XZr+ZjgfTnOr4NJcY1r3f8f0eePS7MRGp+pg==", + "hasInstallScript": true, + "license": "ISC", + "dependencies": { + "prebuild-install": "^7.1.3" + }, + "engines": { + "node": ">=22.0.0" + } + }, "node_modules/js-tokens": { "version": "9.0.1", "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-9.0.1.tgz", @@ -1709,6 +1873,33 @@ "node": ">= 0.6" } }, + "node_modules/mimic-response": { + "version": "3.1.0", + "resolved": "https://registry.npmjs.org/mimic-response/-/mimic-response-3.1.0.tgz", + "integrity": "sha512-z0yWI+4FDrrweS8Zmt4Ej5HdJmky15+L2e6Wgn3+iK5fWzb6T3fhNFq2+MeTRb064c6Wr4N/wv0DzQTjNzHNGQ==", + "license": "MIT", + "engines": { + "node": ">=10" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, + "node_modules/minimist": { + "version": "1.2.8", + "resolved": "https://registry.npmjs.org/minimist/-/minimist-1.2.8.tgz", + "integrity": "sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/mkdirp-classic": { + "version": "0.5.3", + "resolved": "https://registry.npmjs.org/mkdirp-classic/-/mkdirp-classic-0.5.3.tgz", + "integrity": "sha512-gKLcREMhtuZRwRAfqP3RFW+TK4JqApVBtOIftVgjuABpAtpxhPGaDcfvbhNvD0B8iD1oUr/txX35NjcaY6Ns/A==", + "license": "MIT" + }, "node_modules/ms": { "version": "2.1.3", "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", @@ -1734,6 +1925,12 @@ "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1" } }, + "node_modules/napi-build-utils": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/napi-build-utils/-/napi-build-utils-2.0.0.tgz", + "integrity": "sha512-GEbrYkbfF7MoNaoh2iGG84Mnf/WZfB0GdGEsM8wz7Expx/LlWf5U8t9nvJKXSp3qr5IsEbK04cBGhol/KwOsWA==", + "license": "MIT" + }, "node_modules/negotiator": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-1.0.0.tgz", @@ -1743,6 +1940,18 @@ "node": ">= 0.6" } }, + "node_modules/node-abi": { + "version": "3.85.0", + "resolved": "https://registry.npmjs.org/node-abi/-/node-abi-3.85.0.tgz", + "integrity": "sha512-zsFhmbkAzwhTft6nd3VxcG0cvJsT70rL+BIGHWVq5fi6MwGrHwzqKaxXE+Hl2GmnGItnDKPPkO5/LQqjVkIdFg==", + "license": "MIT", + "dependencies": { + "semver": "^7.3.5" + }, + "engines": { + "node": ">=10" + } + }, "node_modules/object-assign": { "version": "4.1.1", "resolved": "https://registry.npmjs.org/object-assign/-/object-assign-4.1.1.tgz", @@ -1888,6 +2097,32 @@ "node": "^10 || ^12 || >=14" } }, + "node_modules/prebuild-install": { + "version": "7.1.3", + "resolved": "https://registry.npmjs.org/prebuild-install/-/prebuild-install-7.1.3.tgz", + "integrity": "sha512-8Mf2cbV7x1cXPUILADGI3wuhfqWvtiLA1iclTDbFRZkgRQS0NqsPZphna9V+HyTEadheuPmjaJMsbzKQFOzLug==", + "license": "MIT", + "dependencies": { + "detect-libc": "^2.0.0", + "expand-template": "^2.0.3", + "github-from-package": "0.0.0", + "minimist": "^1.2.3", + "mkdirp-classic": "^0.5.3", + "napi-build-utils": "^2.0.0", + "node-abi": "^3.3.0", + "pump": "^3.0.0", + "rc": "^1.2.7", + "simple-get": "^4.0.0", + "tar-fs": "^2.0.0", + "tunnel-agent": "^0.6.0" + }, + "bin": { + "prebuild-install": "bin.js" + }, + "engines": { + "node": ">=10" + } + }, "node_modules/proxy-addr": { "version": "2.0.7", "resolved": "https://registry.npmjs.org/proxy-addr/-/proxy-addr-2.0.7.tgz", @@ -1901,6 +2136,16 @@ "node": ">= 0.10" } }, + "node_modules/pump": { + "version": "3.0.3", + "resolved": "https://registry.npmjs.org/pump/-/pump-3.0.3.tgz", + "integrity": "sha512-todwxLMY7/heScKmntwQG8CXVkWUOdYxIvY2s0VWAAMh/nd8SoYiRaKjlr7+iCs984f2P8zvrfWcDDYVb73NfA==", + "license": "MIT", + "dependencies": { + "end-of-stream": "^1.1.0", + "once": "^1.3.1" + } + }, "node_modules/punycode": { "version": "2.3.1", "resolved": "https://registry.npmjs.org/punycode/-/punycode-2.3.1.tgz", @@ -1965,6 +2210,35 @@ "url": "https://opencollective.com/express" } }, + "node_modules/rc": { + "version": "1.2.8", + "resolved": "https://registry.npmjs.org/rc/-/rc-1.2.8.tgz", + "integrity": "sha512-y3bGgqKj3QBdxLbLkomlohkvsA8gdAiUQlSBJnBhfn+BPxg4bc62d8TcBW15wavDfgexCgccckhcZvywyQYPOw==", + "license": "(BSD-2-Clause OR MIT OR Apache-2.0)", + "dependencies": { + "deep-extend": "^0.6.0", + "ini": "~1.3.0", + "minimist": "^1.2.0", + "strip-json-comments": "~2.0.1" + }, + "bin": { + "rc": "cli.js" + } + }, + "node_modules/readable-stream": { + "version": "3.6.2", + "resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-3.6.2.tgz", + "integrity": "sha512-9u/sniCrY3D5WdsERHzHE4G2YCXqoG5FTHUiCC4SIbr6XcLZBY05ya9EKjYek9O5xOAwjGq+1JdGBAS7Q9ScoA==", + "license": "MIT", + "dependencies": { + "inherits": "^2.0.3", + "string_decoder": "^1.1.1", + "util-deprecate": "^1.0.1" + }, + "engines": { + "node": ">= 6" + } + }, "node_modules/rollup": { "version": "4.52.5", "resolved": "https://registry.npmjs.org/rollup/-/rollup-4.52.5.tgz", @@ -2049,6 +2323,18 @@ "integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==", "license": "MIT" }, + "node_modules/semver": { + "version": "7.7.3", + "resolved": "https://registry.npmjs.org/semver/-/semver-7.7.3.tgz", + "integrity": "sha512-SdsKMrI9TdgjdweUSR9MweHA4EJ8YxHn8DFaDisvhVlUOe4BF1tLD7GAj0lIqWVl+dPb/rExr0Btby5loQm20Q==", + "license": "ISC", + "bin": { + "semver": "bin/semver.js" + }, + "engines": { + "node": ">=10" + } + }, "node_modules/send": { "version": "1.2.0", "resolved": "https://registry.npmjs.org/send/-/send-1.2.0.tgz", @@ -2192,6 +2478,51 @@ "dev": true, "license": "ISC" }, + "node_modules/simple-concat": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/simple-concat/-/simple-concat-1.0.1.tgz", + "integrity": "sha512-cSFtAPtRhljv69IK0hTVZQ+OfE9nePi/rtJmw5UjHeVyVroEqJXP1sFztKUy1qU+xvz3u/sfYJLa947b7nAN2Q==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "MIT" + }, + "node_modules/simple-get": { + "version": "4.0.1", + "resolved": "https://registry.npmjs.org/simple-get/-/simple-get-4.0.1.tgz", + "integrity": "sha512-brv7p5WgH0jmQJr1ZDDfKDOSeWWg+OVypG99A/5vYGPqJ6pxiaHLy8nxtFjBA7oMa01ebA9gfh1uMCFqOuXxvA==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "MIT", + "dependencies": { + "decompress-response": "^6.0.0", + "once": "^1.3.1", + "simple-concat": "^1.0.0" + } + }, "node_modules/source-map-js": { "version": "1.2.1", "resolved": "https://registry.npmjs.org/source-map-js/-/source-map-js-1.2.1.tgz", @@ -2225,6 +2556,24 @@ "dev": true, "license": "MIT" }, + "node_modules/string_decoder": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/string_decoder/-/string_decoder-1.3.0.tgz", + "integrity": "sha512-hkRX8U1WjJFd8LsDJ2yQ/wWWxaopEsABU1XfkM8A+j0+85JAGppt16cr1Whg6KIbb4okU6Mql6BOj+uup/wKeA==", + "license": "MIT", + "dependencies": { + "safe-buffer": "~5.2.0" + } + }, + "node_modules/strip-json-comments": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/strip-json-comments/-/strip-json-comments-2.0.1.tgz", + "integrity": "sha512-4gB8na07fecVVkOI6Rs4e7T6NOTki5EmL7TUduTs6bu3EdnSycntVJ4re8kgZA+wx9IueI2Y11bfbgwtzuE0KQ==", + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, "node_modules/strip-literal": { "version": "3.1.0", "resolved": "https://registry.npmjs.org/strip-literal/-/strip-literal-3.1.0.tgz", @@ -2238,6 +2587,34 @@ "url": "https://github.com/sponsors/antfu" } }, + "node_modules/tar-fs": { + "version": "2.1.4", + "resolved": "https://registry.npmjs.org/tar-fs/-/tar-fs-2.1.4.tgz", + "integrity": "sha512-mDAjwmZdh7LTT6pNleZ05Yt65HC3E+NiQzl672vQG38jIrehtJk/J3mNwIg+vShQPcLF/LV7CMnDW6vjj6sfYQ==", + "license": "MIT", + "dependencies": { + "chownr": "^1.1.1", + "mkdirp-classic": "^0.5.2", + "pump": "^3.0.0", + "tar-stream": "^2.1.4" + } + }, + "node_modules/tar-stream": { + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/tar-stream/-/tar-stream-2.2.0.tgz", + "integrity": "sha512-ujeqbceABgwMZxEJnk2HDY2DlnUZ+9oEcb1KzTVfYHio0UE6dG71n60d8D2I4qNvleWrrXpmjpt7vZeF1LnMZQ==", + "license": "MIT", + "dependencies": { + "bl": "^4.0.3", + "end-of-stream": "^1.4.1", + "fs-constants": "^1.0.0", + "inherits": "^2.0.3", + "readable-stream": "^3.1.1" + }, + "engines": { + "node": ">=6" + } + }, "node_modules/tinybench": { "version": "2.9.0", "resolved": "https://registry.npmjs.org/tinybench/-/tinybench-2.9.0.tgz", @@ -2308,6 +2685,18 @@ "node": ">=0.6" } }, + "node_modules/tunnel-agent": { + "version": "0.6.0", + "resolved": "https://registry.npmjs.org/tunnel-agent/-/tunnel-agent-0.6.0.tgz", + "integrity": "sha512-McnNiV1l8RYeY8tBgEpuodCC1mLUdbSN+CYBL7kJsJNInOP8UjDDEwdk6Mw60vdLLrr5NHKZhMAOSrR2NZuQ+w==", + "license": "Apache-2.0", + "dependencies": { + "safe-buffer": "^5.0.1" + }, + "engines": { + "node": "*" + } + }, "node_modules/type-is": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/type-is/-/type-is-2.0.1.tgz", @@ -2361,6 +2750,12 @@ "punycode": "^2.1.0" } }, + "node_modules/util-deprecate": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", + "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==", + "license": "MIT" + }, "node_modules/vary": { "version": "1.1.2", "resolved": "https://registry.npmjs.org/vary/-/vary-1.1.2.tgz", diff --git a/mcp-server/package.json b/mcp-server/package.json index bc68a4b..97b14c8 100644 --- a/mcp-server/package.json +++ b/mcp-server/package.json @@ -17,6 +17,7 @@ }, "dependencies": { "@modelcontextprotocol/sdk": "^1.18.0", + "isolated-vm": "^6.0.2", "zod": "^3.24.2" } } diff --git a/mcp-server/src/poc.test.ts b/mcp-server/src/poc.test.ts index 8ba9cf1..4026e51 100644 --- a/mcp-server/src/poc.test.ts +++ b/mcp-server/src/poc.test.ts @@ -41,48 +41,22 @@ describe('validatePocParams', () => { 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 } + { code: 'console.log("test")' } ); - 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: '' }) + JSON.stringify({ stdout: 'test', 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 } + { code: 'throw new Error("Execution failed")' } ); expect(result.isError).toBe(true); - expect(result.content[0].text).toBe( - JSON.stringify({ error: 'Execution failed' }) - ); + expect(result.content[0].text).toContain('Execution failed'); }); + }); diff --git a/mcp-server/src/poc.ts b/mcp-server/src/poc.ts index bd6b888..4933d94 100644 --- a/mcp-server/src/poc.ts +++ b/mcp-server/src/poc.ts @@ -9,6 +9,7 @@ import { promises as fs } from 'fs'; import path from 'path'; import { exec } from 'child_process'; import { promisify } from 'util'; +import ivm from 'isolated-vm'; const execAsync = promisify(exec); @@ -60,37 +61,38 @@ export async function runPoc( code, }: { code: string; - }, - dependencies: { fs: typeof fs; path: typeof path; execAsync: typeof execAsync } = { fs, path, execAsync } + } ): Promise { 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 isolate = new ivm.Isolate({ memoryLimit: 128 }); + const context = await isolate.createContext(); + const jail = context.global; + await jail.set('global', jail.derefInto()); - const pocFilePath = dependencies.path.join(securityDir, 'poc.js'); - await dependencies.fs.writeFile(pocFilePath, code, 'utf-8'); + const logs: string[] = []; + await jail.set('print', new ivm.Reference((...args: any[]) => { + logs.push(args.map(arg => String(arg)).join(' ')); + })); + await context.eval(` + global.console = { + log: function(...args) { + print.apply(undefined, args, { arguments: { copy: true } }); + }, + error: function(...args) { + print.apply(undefined, args, { arguments: { copy: true } }); + } + }; + `); - 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}`); + const script = await isolate.compileScript(code); + await script.run(context); return { content: [ { type: 'text', - text: JSON.stringify({ stdout, stderr }), + text: JSON.stringify({ stdout: logs.join('\n'), stderr: '' }), }, ], };