From b62abfea4c675f83ee3369993add1b78d84bc056 Mon Sep 17 00:00:00 2001 From: ezl-keygraph Date: Tue, 3 Mar 2026 01:08:26 +0530 Subject: [PATCH 1/3] feat: add three-tier model system with Bedrock support Introduce small/medium/large model tiers so agents use the appropriate model for their task complexity. Pre-recon uses Opus (large) for deep source code analysis, most agents use Sonnet (medium), and report uses Haiku (small) for summarization. - Add src/ai/models.ts with ModelTier type and resolveModel() - Add modelTier field to AgentDefinition - Refactor claude-executor env var passthrough into loop - Add Bedrock credential validation in preflight and CLI - Pass through Bedrock and model env vars in docker-compose --- .env.example | 22 ++++++++++++++++++++ README.md | 29 ++++++++++++++++++++++++++ docker-compose.yml | 6 ++++++ shannon | 19 ++++++++++++++--- src/ai/claude-executor.ts | 33 +++++++++++++++++------------ src/ai/models.ts | 37 +++++++++++++++++++++++++++++++++ src/services/agent-execution.ts | 3 ++- src/services/preflight.ts | 28 +++++++++++++++++++++---- src/session-manager.ts | 2 ++ src/types/agents.ts | 1 + 10 files changed, 159 insertions(+), 21 deletions(-) create mode 100644 src/ai/models.ts diff --git a/.env.example b/.env.example index b9aab98..d5d101e 100644 --- a/.env.example +++ b/.env.example @@ -26,6 +26,28 @@ ANTHROPIC_API_KEY=your-api-key-here # OPENROUTER_API_KEY=sk-or-your-openrouter-key # ROUTER_DEFAULT=openrouter,google/gemini-3-flash-preview +# ============================================================================= +# Model Tier Overrides (Anthropic API / OAuth / Bedrock) +# ============================================================================= +# Override which model is used for each tier. Defaults are used if not set. +# ANTHROPIC_SMALL_MODEL=... # Small tier (default: claude-haiku-4-5-20251001) +# ANTHROPIC_MEDIUM_MODEL=... # Medium tier (default: claude-sonnet-4-6) +# ANTHROPIC_LARGE_MODEL=... # Large tier (default: claude-opus-4-6) + +# ============================================================================= +# OPTION 3: AWS Bedrock +# ============================================================================= +# https://aws.amazon.com/blogs/machine-learning/accelerate-ai-development-with-amazon-bedrock-api-keys/ +# Requires the model tier overrides above to be set with Bedrock-specific model IDs. +# Example Bedrock model IDs for us-east-1: +# ANTHROPIC_SMALL_MODEL=us.anthropic.claude-haiku-4-5-20251001-v1:0 +# ANTHROPIC_MEDIUM_MODEL=us.anthropic.claude-sonnet-4-6 +# ANTHROPIC_LARGE_MODEL=us.anthropic.claude-opus-4-6 + +# CLAUDE_CODE_USE_BEDROCK=1 +# AWS_REGION=us-east-1 +# AWS_BEARER_TOKEN_BEDROCK=your-bearer-token + # ============================================================================= # Available Models # ============================================================================= diff --git a/README.md b/README.md index dbf35a7..74c4a25 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,7 @@ Shannon is available in two editions: - [Usage Examples](#usage-examples) - [Workspaces and Resuming](#workspaces-and-resuming) - [Configuration (Optional)](#configuration-optional) + - [AWS Bedrock](#aws-bedrock) - [[EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers)](#experimental---unsupported-router-mode-alternative-providers) - [Output and Results](#output-and-results) - [Sample Reports](#-sample-reports) @@ -107,6 +108,7 @@ Shannon is available in two editions: - **AI Provider Credentials** (choose one): - **Anthropic API key** (recommended) - Get from [Anthropic Console](https://console.anthropic.com) - **Claude Code OAuth token** + - **AWS Bedrock** - Route through Amazon Bedrock with AWS credentials (see [AWS Bedrock](#aws-bedrock)) - **[EXPERIMENTAL - UNSUPPORTED] Alternative providers via Router Mode** - OpenAI or Google Gemini via OpenRouter (see [Router Mode](#experimental---unsupported-router-mode-alternative-providers)) ### Quick Start @@ -348,6 +350,33 @@ pipeline: `max_concurrent_pipelines` controls how many vulnerability pipelines run simultaneously (1-5, default: 5). Lower values reduce the chance of hitting rate limits but increase wall-clock time. +### AWS Bedrock + +Shannon also supports [Amazon Bedrock](https://aws.amazon.com/bedrock/) instead of using an Anthropic API key. + +#### Quick Setup + +1. Add your AWS credentials to `.env`: + +```bash +CLAUDE_CODE_USE_BEDROCK=1 +AWS_REGION=us-east-1 +AWS_BEARER_TOKEN_BEDROCK=your-bearer-token + +# Set models with Bedrock-specific IDs for your region +ANTHROPIC_SMALL_MODEL=us.anthropic.claude-haiku-4-5-20251001-v1:0 +ANTHROPIC_MEDIUM_MODEL=us.anthropic.claude-sonnet-4-6 +ANTHROPIC_LARGE_MODEL=us.anthropic.claude-opus-4-6 +``` + +2. Run Shannon as usual: + +```bash +./shannon start URL=https://example.com REPO=repo-name +``` + +Shannon uses three model tiers: **small** (`claude-haiku-4-5-20251001`) for summarization, **medium** (`claude-sonnet-4-6`) for security analysis, and **large** (`claude-opus-4-6`) for deep reasoning. Set `ANTHROPIC_SMALL_MODEL`, `ANTHROPIC_MEDIUM_MODEL`, and `ANTHROPIC_LARGE_MODEL` to the Bedrock model IDs for your region. + ### [EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers) Shannon can experimentally route requests through alternative AI providers using claude-code-router. This mode is not officially supported and is intended primarily for: diff --git a/docker-compose.yml b/docker-compose.yml index e54ba1f..591b075 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,6 +24,12 @@ services: - ANTHROPIC_AUTH_TOKEN=${ANTHROPIC_AUTH_TOKEN:-} # Auth token for router - ROUTER_DEFAULT=${ROUTER_DEFAULT:-} # Model name when using router (e.g., "gemini,gemini-2.5-pro") - CLAUDE_CODE_OAUTH_TOKEN=${CLAUDE_CODE_OAUTH_TOKEN:-} + - CLAUDE_CODE_USE_BEDROCK=${CLAUDE_CODE_USE_BEDROCK:-} + - AWS_REGION=${AWS_REGION:-} + - AWS_BEARER_TOKEN_BEDROCK=${AWS_BEARER_TOKEN_BEDROCK:-} + - ANTHROPIC_SMALL_MODEL=${ANTHROPIC_SMALL_MODEL:-} + - ANTHROPIC_MEDIUM_MODEL=${ANTHROPIC_MEDIUM_MODEL:-} + - ANTHROPIC_LARGE_MODEL=${ANTHROPIC_LARGE_MODEL:-} - CLAUDE_CODE_MAX_OUTPUT_TOKENS=${CLAUDE_CODE_MAX_OUTPUT_TOKENS:-64000} depends_on: temporal: diff --git a/shannon b/shannon index 70a74e0..8606d3b 100755 --- a/shannon +++ b/shannon @@ -142,14 +142,27 @@ cmd_start() { exit 1 fi - # Check for API key (router mode can use alternative provider API keys) + # Check for API key (Bedrock and router modes can bypass this) if [ -z "$ANTHROPIC_API_KEY" ] && [ -z "$CLAUDE_CODE_OAUTH_TOKEN" ]; then - if [ "$ROUTER" = "true" ] && { [ -n "$OPENAI_API_KEY" ] || [ -n "$OPENROUTER_API_KEY" ]; }; then + if [ "$CLAUDE_CODE_USE_BEDROCK" = "1" ]; then + # Bedrock mode — validate required AWS credentials + MISSING="" + [ -z "$AWS_REGION" ] && MISSING="$MISSING AWS_REGION" + [ -z "$AWS_BEARER_TOKEN_BEDROCK" ] && MISSING="$MISSING AWS_BEARER_TOKEN_BEDROCK" + [ -z "$ANTHROPIC_SMALL_MODEL" ] && MISSING="$MISSING ANTHROPIC_SMALL_MODEL" + [ -z "$ANTHROPIC_MEDIUM_MODEL" ] && MISSING="$MISSING ANTHROPIC_MEDIUM_MODEL" + [ -z "$ANTHROPIC_LARGE_MODEL" ] && MISSING="$MISSING ANTHROPIC_LARGE_MODEL" + if [ -n "$MISSING" ]; then + echo "ERROR: Bedrock mode requires the following env vars in .env:$MISSING" + exit 1 + fi + elif [ "$ROUTER" = "true" ] && { [ -n "$OPENAI_API_KEY" ] || [ -n "$OPENROUTER_API_KEY" ]; }; then # Router mode with alternative provider - set a placeholder for SDK init export ANTHROPIC_API_KEY="router-mode" else echo "ERROR: Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env" - echo " (or use ROUTER=true with OPENAI_API_KEY or OPENROUTER_API_KEY)" + echo " (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock," + echo " or ROUTER=true with OPENAI_API_KEY or OPENROUTER_API_KEY)" exit 1 fi fi diff --git a/src/ai/claude-executor.ts b/src/ai/claude-executor.ts index 921c9a7..6a3af0c 100644 --- a/src/ai/claude-executor.ts +++ b/src/ai/claude-executor.ts @@ -24,6 +24,7 @@ import { detectExecutionContext, formatErrorOutput, formatCompletionMessage } fr import { createProgressManager } from './progress-manager.js'; import { createAuditLogger } from './audit-logger.js'; import { getActualModelName } from './router-utils.js'; +import { resolveModel, type ModelTier } from './models.js'; import type { ActivityLogger } from '../types/activity-logger.js'; declare global { @@ -202,7 +203,8 @@ export async function runClaudePrompt( description: string = 'Claude analysis', agentName: string | null = null, auditSession: AuditSession | null = null, - logger: ActivityLogger + logger: ActivityLogger, + modelTier: ModelTier = 'medium' ): Promise { // 1. Initialize timing and prompt const timer = new Timer(`agent-${description.toLowerCase().replace(/\s+/g, '-')}`); @@ -225,22 +227,27 @@ export async function runClaudePrompt( const sdkEnv: Record = { CLAUDE_CODE_MAX_OUTPUT_TOKENS: process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS || '64000', }; - if (process.env.ANTHROPIC_API_KEY) { - sdkEnv.ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY; - } - if (process.env.CLAUDE_CODE_OAUTH_TOKEN) { - sdkEnv.CLAUDE_CODE_OAUTH_TOKEN = process.env.CLAUDE_CODE_OAUTH_TOKEN; - } - if (process.env.ANTHROPIC_BASE_URL) { - sdkEnv.ANTHROPIC_BASE_URL = process.env.ANTHROPIC_BASE_URL; - } - if (process.env.ANTHROPIC_AUTH_TOKEN) { - sdkEnv.ANTHROPIC_AUTH_TOKEN = process.env.ANTHROPIC_AUTH_TOKEN; + const passthroughVars = [ + 'ANTHROPIC_API_KEY', + 'CLAUDE_CODE_OAUTH_TOKEN', + 'ANTHROPIC_BASE_URL', + 'ANTHROPIC_AUTH_TOKEN', + 'CLAUDE_CODE_USE_BEDROCK', + 'AWS_REGION', + 'AWS_BEARER_TOKEN_BEDROCK', + 'ANTHROPIC_SMALL_MODEL', + 'ANTHROPIC_MEDIUM_MODEL', + 'ANTHROPIC_LARGE_MODEL', + ]; + for (const name of passthroughVars) { + if (process.env[name]) { + sdkEnv[name] = process.env[name]!; + } } // 5. Configure SDK options const options = { - model: 'claude-sonnet-4-5-20250929', + model: resolveModel(modelTier), maxTurns: 10_000, cwd: sourceDir, permissionMode: 'bypassPermissions' as const, diff --git a/src/ai/models.ts b/src/ai/models.ts new file mode 100644 index 0000000..7728215 --- /dev/null +++ b/src/ai/models.ts @@ -0,0 +1,37 @@ +// Copyright (C) 2025 Keygraph, Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License version 3 +// as published by the Free Software Foundation. + +/** + * Model tier definitions and resolution. + * + * Three tiers mapped to capability levels: + * - "small" (Haiku — summarization, structured extraction) + * - "medium" (Sonnet — tool use, general analysis) + * - "large" (Opus — deep reasoning, complex analysis) + * + * Users override via ANTHROPIC_SMALL_MODEL / ANTHROPIC_MEDIUM_MODEL / ANTHROPIC_LARGE_MODEL, + * which works across all providers (direct, Bedrock, Vertex). + */ + +export type ModelTier = 'small' | 'medium' | 'large'; + +const DEFAULT_MODELS: Readonly> = { + small: 'claude-haiku-4-5-20251001', + medium: 'claude-sonnet-4-6', + large: 'claude-opus-4-6', +}; + +/** Resolve a model tier to a concrete model ID. */ +export function resolveModel(tier: ModelTier = 'medium'): string { + switch (tier) { + case 'small': + return process.env.ANTHROPIC_SMALL_MODEL || DEFAULT_MODELS.small; + case 'large': + return process.env.ANTHROPIC_LARGE_MODEL || DEFAULT_MODELS.large; + default: + return process.env.ANTHROPIC_MEDIUM_MODEL || DEFAULT_MODELS.medium; + } +} diff --git a/src/services/agent-execution.ts b/src/services/agent-execution.ts index 401f051..f0f01df 100644 --- a/src/services/agent-execution.ts +++ b/src/services/agent-execution.ts @@ -156,7 +156,8 @@ export class AgentExecutionService { agentName, // description agentName, auditSession, - logger + logger, + AGENTS[agentName].modelTier ); // 6. Spending cap check - defense-in-depth diff --git a/src/services/preflight.ts b/src/services/preflight.ts index 95d6a08..dd2dc15 100644 --- a/src/services/preflight.ts +++ b/src/services/preflight.ts @@ -24,6 +24,7 @@ import { PentestError, isRetryableError } from './error-handling.js'; import { ErrorCode } from '../types/errors.js'; import { type Result, ok, err } from '../types/result.js'; import { parseConfig } from '../config-parser.js'; +import { resolveModel } from '../ai/models.js'; import type { ActivityLogger } from '../types/activity-logger.js'; // === Repository Validation === @@ -165,11 +166,30 @@ async function validateCredentials( return ok(undefined); } - // 2. Check that at least one credential is present + // 2. Bedrock mode — validate required AWS credentials are present + if (process.env.CLAUDE_CODE_USE_BEDROCK === '1') { + const required = ['AWS_REGION', 'AWS_BEARER_TOKEN_BEDROCK', 'ANTHROPIC_SMALL_MODEL', 'ANTHROPIC_MEDIUM_MODEL', 'ANTHROPIC_LARGE_MODEL']; + const missing = required.filter(v => !process.env[v]); + if (missing.length > 0) { + return err( + new PentestError( + `Bedrock mode requires the following env vars in .env: ${missing.join(', ')}`, + 'config', + false, + { missing }, + ErrorCode.AUTH_FAILED + ) + ); + } + logger.info('Bedrock credentials OK'); + return ok(undefined); + } + + // 3. Check that at least one credential is present if (!process.env.ANTHROPIC_API_KEY && !process.env.CLAUDE_CODE_OAUTH_TOKEN) { return err( new PentestError( - 'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env', + 'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock)', 'config', false, {}, @@ -178,12 +198,12 @@ async function validateCredentials( ); } - // 3. Validate via SDK query + // 4. Validate via SDK query const authType = process.env.CLAUDE_CODE_OAUTH_TOKEN ? 'OAuth token' : 'API key'; logger.info(`Validating ${authType} via SDK...`); try { - for await (const message of query({ prompt: 'hi', options: { model: 'claude-haiku-4-5-20251001', maxTurns: 1 } })) { + for await (const message of query({ prompt: 'hi', options: { model: resolveModel('small'), maxTurns: 1 } })) { if (message.type === 'assistant' && message.error) { return classifySdkError(message.error, authType); } diff --git a/src/session-manager.ts b/src/session-manager.ts index 8180040..a17cc7d 100644 --- a/src/session-manager.ts +++ b/src/session-manager.ts @@ -18,6 +18,7 @@ export const AGENTS: Readonly> = Object.freez prerequisites: [], promptTemplate: 'pre-recon-code', deliverableFilename: 'code_analysis_deliverable.md', + modelTier: 'large', }, 'recon': { name: 'recon', @@ -102,6 +103,7 @@ export const AGENTS: Readonly> = Object.freez prerequisites: ['injection-exploit', 'xss-exploit', 'auth-exploit', 'ssrf-exploit', 'authz-exploit'], promptTemplate: 'report-executive', deliverableFilename: 'comprehensive_security_assessment_report.md', + modelTier: 'small', }, }); diff --git a/src/types/agents.ts b/src/types/agents.ts index e0265a9..41b5873 100644 --- a/src/types/agents.ts +++ b/src/types/agents.ts @@ -58,6 +58,7 @@ export interface AgentDefinition { prerequisites: AgentName[]; promptTemplate: string; deliverableFilename: string; + modelTier?: 'small' | 'medium' | 'large'; } /** From 3ec491b30b5f736237c744a8e92ebb6931182d87 Mon Sep 17 00:00:00 2001 From: ezl-keygraph Date: Tue, 3 Mar 2026 02:05:09 +0530 Subject: [PATCH 2/3] chore: update pipeline testing vulnerability prompts --- prompts/pipeline-testing/vuln-auth.txt | 44 +++++---------------- prompts/pipeline-testing/vuln-authz.txt | 44 +++++---------------- prompts/pipeline-testing/vuln-injection.txt | 44 +++++---------------- prompts/pipeline-testing/vuln-ssrf.txt | 44 +++++---------------- prompts/pipeline-testing/vuln-xss.txt | 44 +++++---------------- 5 files changed, 50 insertions(+), 170 deletions(-) diff --git a/prompts/pipeline-testing/vuln-auth.txt b/prompts/pipeline-testing/vuln-auth.txt index afc5116..8ba4fea 100644 --- a/prompts/pipeline-testing/vuln-auth.txt +++ b/prompts/pipeline-testing/vuln-auth.txt @@ -1,37 +1,13 @@ -## 🧪 Pipeline Testing: MCP Isolation Test for Auth Agent +Please complete these tasks using your MCP tools: -**MCP Server Assignment:** Using `{{MCP_SERVER}}` for browser automation testing. +1. Navigate to https://example.net and take a screenshot: + - Use {{MCP_SERVER}}__browser_navigate to go to https://example.net + - Use {{MCP_SERVER}}__browser_take_screenshot -### Step 1: Test Browser Isolation -1. Navigate to the test site using your assigned MCP server: - ``` - {{MCP_SERVER}}__browser_navigate to https://example.net - ``` +2. Save an analysis deliverable: + - Use save_deliverable with deliverable_type "AUTH_ANALYSIS" + - Content: "# Auth Analysis Report\n\nAnalysis complete. No authentication vulnerabilities identified." -2. Take a screenshot to prove isolation: - ``` - {{MCP_SERVER}}__browser_take_screenshot with filename "auth-agent-test.png" - ``` - -3. Log the test results: - ``` - console.log("🔵 AUTH AGENT - MCP Server: {{MCP_SERVER}} - Successfully navigated to example.net") - ``` - -### Step 2: Create Required Deliverables -Use the save_deliverable MCP tool to create both deliverables: - -**1. Analysis Report:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "AUTH_ANALYSIS"` -- `content: "# Auth Analysis Complete (Pipeline Test)\n\n**MCP Server Used:** {{MCP_SERVER}}\n**Test Site:** https://example.net\n**Screenshot:** auth-agent-test.png\n**Status:** Browser isolation test successful"` - -**2. Exploitation Queue:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "AUTH_QUEUE"` -- `content: {"vulnerabilities": [{"ID": "AUTH-VULN-01", "verdict": "vulnerable", "confidence": "high", "notes": "Test auth vulnerability for pipeline validation (MCP: {{MCP_SERVER}})", "vulnerability_type": "Authentication_Bypass", "endpoint": "/api/login", "flaw_type": "Missing token validation", "exploit_hypothesis": "Empty or invalid token allows access to protected resources", "witness_payload": "Authorization: Bearer invalid_token"}]}` - -These tools will automatically create the correct files in `deliverables/`. - -### Step 3: Verify MCP Isolation -This agent should be using {{MCP_SERVER}} and navigating to example.net independently of other parallel agents. \ No newline at end of file +3. Save a queue deliverable: + - Use save_deliverable with deliverable_type "AUTH_QUEUE" + - Content: {"vulnerabilities": []} diff --git a/prompts/pipeline-testing/vuln-authz.txt b/prompts/pipeline-testing/vuln-authz.txt index 5f42cc0..bb005ec 100644 --- a/prompts/pipeline-testing/vuln-authz.txt +++ b/prompts/pipeline-testing/vuln-authz.txt @@ -1,37 +1,13 @@ -## 🧪 Pipeline Testing: MCP Isolation Test for Authorization Agent +Please complete these tasks using your MCP tools: -**MCP Server Assignment:** Using `{{MCP_SERVER}}` for browser automation testing. +1. Navigate to https://jsonplaceholder.typicode.com and take a screenshot: + - Use {{MCP_SERVER}}__browser_navigate to go to https://jsonplaceholder.typicode.com + - Use {{MCP_SERVER}}__browser_take_screenshot -### Step 1: Test Browser Isolation -1. Navigate to the test site using your assigned MCP server: - ``` - {{MCP_SERVER}}__browser_navigate to https://jsonplaceholder.typicode.com - ``` +2. Save an analysis deliverable: + - Use save_deliverable with deliverable_type "AUTHZ_ANALYSIS" + - Content: "# Authorization Analysis Report\n\nAnalysis complete. No authorization vulnerabilities identified." -2. Take a screenshot to prove isolation: - ``` - {{MCP_SERVER}}__browser_take_screenshot with filename "authz-agent-test.png" - ``` - -3. Log the test results: - ``` - console.log("🟢 AUTHZ AGENT - MCP Server: {{MCP_SERVER}} - Successfully navigated to jsonplaceholder.typicode.com") - ``` - -### Step 2: Create Required Deliverables -Use the save_deliverable MCP tool to create both deliverables: - -**1. Analysis Report:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "AUTHZ_ANALYSIS"` -- `content: "# Authorization Analysis Complete (Pipeline Test)\n\n**MCP Server Used:** {{MCP_SERVER}}\n**Test Site:** https://jsonplaceholder.typicode.com\n**Screenshot:** authz-agent-test.png\n**Status:** Browser isolation test successful"` - -**2. Exploitation Queue:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "AUTHZ_QUEUE"` -- `content: {"vulnerabilities": [{"ID": "AUTHZ-VULN-01", "verdict": "vulnerable", "confidence": "high", "notes": "Test authz vulnerability for pipeline validation (MCP: {{MCP_SERVER}})", "vulnerability_type": "Vertical", "endpoint": "/admin/users", "actual_access": "Regular users can access admin functions", "witness_payload": "GET /admin/users with regular user token"}]}` - -These tools will automatically create the correct files in `deliverables/`. - -### Step 3: Verify MCP Isolation -This agent should be using {{MCP_SERVER}} and navigating to jsonplaceholder.typicode.com independently of other parallel agents. \ No newline at end of file +3. Save a queue deliverable: + - Use save_deliverable with deliverable_type "AUTHZ_QUEUE" + - Content: {"vulnerabilities": []} diff --git a/prompts/pipeline-testing/vuln-injection.txt b/prompts/pipeline-testing/vuln-injection.txt index 9b0c842..66a822a 100644 --- a/prompts/pipeline-testing/vuln-injection.txt +++ b/prompts/pipeline-testing/vuln-injection.txt @@ -1,37 +1,13 @@ -## 🧪 Pipeline Testing: MCP Isolation Test for Injection Agent +Please complete these tasks using your MCP tools: -**MCP Server Assignment:** Using `{{MCP_SERVER}}` for browser automation testing. +1. Navigate to https://example.com and take a screenshot: + - Use {{MCP_SERVER}}__browser_navigate to go to https://example.com + - Use {{MCP_SERVER}}__browser_take_screenshot -### Step 1: Test Browser Isolation -1. Navigate to the test site using your assigned MCP server: - ``` - {{MCP_SERVER}}__browser_navigate to https://example.com - ``` +2. Save an analysis deliverable: + - Use save_deliverable with deliverable_type "INJECTION_ANALYSIS" + - Content: "# Injection Analysis Report\n\nAnalysis complete. No injection vulnerabilities identified." -2. Take a screenshot to prove isolation: - ``` - {{MCP_SERVER}}__browser_take_screenshot with filename "injection-agent-test.png" - ``` - -3. Log the test results: - ``` - console.log("🔴 INJECTION AGENT - MCP Server: {{MCP_SERVER}} - Successfully navigated to example.com") - ``` - -### Step 2: Create Required Deliverables -Use the save_deliverable MCP tool to create both deliverables: - -**1. Analysis Report:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "INJECTION_ANALYSIS"` -- `content: "# Injection Analysis Complete (Pipeline Test)\n\n**MCP Server Used:** {{MCP_SERVER}}\n**Test Site:** https://example.com\n**Screenshot:** injection-agent-test.png\n**Status:** Browser isolation test successful"` - -**2. Exploitation Queue:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "INJECTION_QUEUE"` -- `content: {"vulnerabilities": [{"ID": "INJ-VULN-01", "verdict": "vulnerable", "confidence": "high", "notes": "Test vulnerability for pipeline validation (MCP: {{MCP_SERVER}})", "vulnerability_type": "SQLi", "source": "query parameter at app.js:15", "path": "controller → db.query", "sink_call": "db.query at app.js:20", "slot_type": "SQL-val", "mismatch_reason": "User input directly concatenated into SQL query", "witness_payload": "'; DROP TABLE users; --"}]}` - -These tools will automatically create the correct files in `deliverables/`. - -### Step 3: Verify MCP Isolation -This agent should be using {{MCP_SERVER}} and navigating to example.com independently of other parallel agents. \ No newline at end of file +3. Save a queue deliverable: + - Use save_deliverable with deliverable_type "INJECTION_QUEUE" + - Content: {"vulnerabilities": []} diff --git a/prompts/pipeline-testing/vuln-ssrf.txt b/prompts/pipeline-testing/vuln-ssrf.txt index 9198edd..9662cab 100644 --- a/prompts/pipeline-testing/vuln-ssrf.txt +++ b/prompts/pipeline-testing/vuln-ssrf.txt @@ -1,37 +1,13 @@ -## 🧪 Pipeline Testing: MCP Isolation Test for SSRF Agent +Please complete these tasks using your MCP tools: -**MCP Server Assignment:** Using `{{MCP_SERVER}}` for browser automation testing. +1. Navigate to https://httpbin.org and take a screenshot: + - Use {{MCP_SERVER}}__browser_navigate to go to https://httpbin.org + - Use {{MCP_SERVER}}__browser_take_screenshot -### Step 1: Test Browser Isolation -1. Navigate to the test site using your assigned MCP server: - ``` - {{MCP_SERVER}}__browser_navigate to https://httpbin.org - ``` +2. Save an analysis deliverable: + - Use save_deliverable with deliverable_type "SSRF_ANALYSIS" + - Content: "# SSRF Analysis Report\n\nAnalysis complete. No SSRF vulnerabilities identified." -2. Take a screenshot to prove isolation: - ``` - {{MCP_SERVER}}__browser_take_screenshot with filename "ssrf-agent-test.png" - ``` - -3. Log the test results: - ``` - console.log("🟠 SSRF AGENT - MCP Server: {{MCP_SERVER}} - Successfully navigated to httpbin.org") - ``` - -### Step 2: Create Required Deliverables -Use the save_deliverable MCP tool to create both deliverables: - -**1. Analysis Report:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "SSRF_ANALYSIS"` -- `content: "# SSRF Analysis Complete (Pipeline Test)\n\n**MCP Server Used:** {{MCP_SERVER}}\n**Test Site:** https://httpbin.org\n**Screenshot:** ssrf-agent-test.png\n**Status:** Browser isolation test successful"` - -**2. Exploitation Queue:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "SSRF_QUEUE"` -- `content: {"vulnerabilities": [{"ID": "SSRF-VULN-01", "verdict": "vulnerable", "confidence": "high", "notes": "Test SSRF vulnerability for pipeline validation (MCP: {{MCP_SERVER}})", "vulnerability_type": "URL_Manipulation", "source": "url parameter in /api/fetch", "outbound_call": "fetch() at api.js:45", "witness_payload": "http://internal.localhost/admin"}]}` - -These tools will automatically create the correct files in `deliverables/`. - -### Step 3: Verify MCP Isolation -This agent should be using {{MCP_SERVER}} and navigating to httpbin.org independently of other parallel agents. \ No newline at end of file +3. Save a queue deliverable: + - Use save_deliverable with deliverable_type "SSRF_QUEUE" + - Content: {"vulnerabilities": []} diff --git a/prompts/pipeline-testing/vuln-xss.txt b/prompts/pipeline-testing/vuln-xss.txt index 23c4f0e..d048245 100644 --- a/prompts/pipeline-testing/vuln-xss.txt +++ b/prompts/pipeline-testing/vuln-xss.txt @@ -1,37 +1,13 @@ -## 🧪 Pipeline Testing: MCP Isolation Test for XSS Agent +Please complete these tasks using your MCP tools: -**MCP Server Assignment:** Using `{{MCP_SERVER}}` for browser automation testing. +1. Navigate to https://example.org and take a screenshot: + - Use {{MCP_SERVER}}__browser_navigate to go to https://example.org + - Use {{MCP_SERVER}}__browser_take_screenshot -### Step 1: Test Browser Isolation -1. Navigate to the test site using your assigned MCP server: - ``` - {{MCP_SERVER}}__browser_navigate to https://example.org - ``` +2. Save an analysis deliverable: + - Use save_deliverable with deliverable_type "XSS_ANALYSIS" + - Content: "# XSS Analysis Report\n\nAnalysis complete. No XSS vulnerabilities identified." -2. Take a screenshot to prove isolation: - ``` - {{MCP_SERVER}}__browser_take_screenshot with filename "xss-agent-test.png" - ``` - -3. Log the test results: - ``` - console.log("🟡 XSS AGENT - MCP Server: {{MCP_SERVER}} - Successfully navigated to example.org") - ``` - -### Step 2: Create Required Deliverables -Use the save_deliverable MCP tool to create both deliverables: - -**1. Analysis Report:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "XSS_ANALYSIS"` -- `content: "# XSS Analysis Complete (Pipeline Test)\n\n**MCP Server Used:** {{MCP_SERVER}}\n**Test Site:** https://example.org\n**Screenshot:** xss-agent-test.png\n**Status:** Browser isolation test successful"` - -**2. Exploitation Queue:** -Use `save_deliverable` MCP tool with: -- `deliverable_type: "XSS_QUEUE"` -- `content: {"vulnerabilities": [{"ID": "XSS-VULN-01", "verdict": "vulnerable", "confidence": "high", "notes": "Test XSS vulnerability for pipeline validation (MCP: {{MCP_SERVER}})", "vulnerability_type": "Reflected", "source": "search parameter", "sink_function": "template.render at search.js:25", "render_context": "HTML_BODY", "mismatch_reason": "User input rendered without HTML encoding", "witness_payload": ""}]}` - -These tools will automatically create the correct files in `deliverables/`. - -### Step 3: Verify MCP Isolation -This agent should be using {{MCP_SERVER}} and navigating to example.org independently of other parallel agents. \ No newline at end of file +3. Save a queue deliverable: + - Use save_deliverable with deliverable_type "XSS_QUEUE" + - Content: {"vulnerabilities": []} From 6a76df2f4c653778c67dfc375ff07a84696ffcfe Mon Sep 17 00:00:00 2001 From: ezl-keygraph Date: Tue, 3 Mar 2026 02:42:46 +0530 Subject: [PATCH 3/3] feat: add Google Vertex AI support with service account auth --- .env.example | 17 +++++++++++++ .gitignore | 1 + README.md | 39 ++++++++++++++++++++++++++++ docker-compose.yml | 5 ++++ shannon | 27 +++++++++++++++++++- src/ai/claude-executor.ts | 4 +++ src/services/preflight.ts | 53 ++++++++++++++++++++++++++++++++++++--- 7 files changed, 141 insertions(+), 5 deletions(-) diff --git a/.env.example b/.env.example index d5d101e..4c01091 100644 --- a/.env.example +++ b/.env.example @@ -48,6 +48,23 @@ ANTHROPIC_API_KEY=your-api-key-here # AWS_REGION=us-east-1 # AWS_BEARER_TOKEN_BEDROCK=your-bearer-token +# ============================================================================= +# OPTION 4: Google Vertex AI +# ============================================================================= +# https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-partner-models +# Requires a GCP service account with roles/aiplatform.user. +# Download the SA key JSON from GCP Console (IAM > Service Accounts > Keys). +# Requires the model tier overrides above to be set with Vertex AI model IDs. +# Example Vertex AI model IDs: +# ANTHROPIC_SMALL_MODEL=claude-haiku-4-5@20251001 +# ANTHROPIC_MEDIUM_MODEL=claude-sonnet-4-6 +# ANTHROPIC_LARGE_MODEL=claude-opus-4-6 + +# CLAUDE_CODE_USE_VERTEX=1 +# CLOUD_ML_REGION=us-east5 +# ANTHROPIC_VERTEX_PROJECT_ID=your-gcp-project-id +# GOOGLE_APPLICATION_CREDENTIALS=./credentials/gcp-sa-key.json + # ============================================================================= # Available Models # ============================================================================= diff --git a/.gitignore b/.gitignore index 785ad38..1c1aa3a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ node_modules/ .env audit-logs/ +credentials/ dist/ repos/ diff --git a/README.md b/README.md index 74c4a25..21b2a0a 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,7 @@ Shannon is available in two editions: - [Workspaces and Resuming](#workspaces-and-resuming) - [Configuration (Optional)](#configuration-optional) - [AWS Bedrock](#aws-bedrock) + - [Google Vertex AI](#google-vertex-ai) - [[EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers)](#experimental---unsupported-router-mode-alternative-providers) - [Output and Results](#output-and-results) - [Sample Reports](#-sample-reports) @@ -109,6 +110,7 @@ Shannon is available in two editions: - **Anthropic API key** (recommended) - Get from [Anthropic Console](https://console.anthropic.com) - **Claude Code OAuth token** - **AWS Bedrock** - Route through Amazon Bedrock with AWS credentials (see [AWS Bedrock](#aws-bedrock)) + - **Google Vertex AI** - Route through Google Cloud Vertex AI (see [Google Vertex AI](#google-vertex-ai)) - **[EXPERIMENTAL - UNSUPPORTED] Alternative providers via Router Mode** - OpenAI or Google Gemini via OpenRouter (see [Router Mode](#experimental---unsupported-router-mode-alternative-providers)) ### Quick Start @@ -377,6 +379,43 @@ ANTHROPIC_LARGE_MODEL=us.anthropic.claude-opus-4-6 Shannon uses three model tiers: **small** (`claude-haiku-4-5-20251001`) for summarization, **medium** (`claude-sonnet-4-6`) for security analysis, and **large** (`claude-opus-4-6`) for deep reasoning. Set `ANTHROPIC_SMALL_MODEL`, `ANTHROPIC_MEDIUM_MODEL`, and `ANTHROPIC_LARGE_MODEL` to the Bedrock model IDs for your region. +### Google Vertex AI + +Shannon also supports [Google Vertex AI](https://cloud.google.com/vertex-ai) instead of using an Anthropic API key. + +#### Quick Setup + +1. Create a service account with the `roles/aiplatform.user` role in the [GCP Console](https://console.cloud.google.com/iam-admin/serviceaccounts), then download a JSON key file. + +2. Place the key file in the `./credentials/` directory: + +```bash +mkdir -p ./credentials +cp /path/to/your-sa-key.json ./credentials/gcp-sa-key.json +``` + +3. Add your GCP configuration to `.env`: + +```bash +CLAUDE_CODE_USE_VERTEX=1 +CLOUD_ML_REGION=us-east5 +ANTHROPIC_VERTEX_PROJECT_ID=your-gcp-project-id +GOOGLE_APPLICATION_CREDENTIALS=./credentials/gcp-sa-key.json + +# Set models with Vertex AI model IDs +ANTHROPIC_SMALL_MODEL=claude-haiku-4-5@20251001 +ANTHROPIC_MEDIUM_MODEL=claude-sonnet-4-6 +ANTHROPIC_LARGE_MODEL=claude-opus-4-6 +``` + +4. Run Shannon as usual: + +```bash +./shannon start URL=https://example.com REPO=repo-name +``` + +Set `CLOUD_ML_REGION=global` for global endpoints, or a specific region like `us-east5`. Some models may not be available on global endpoints — see the [Vertex AI Model Garden](https://console.cloud.google.com/vertex-ai/model-garden) for region availability. + ### [EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers) Shannon can experimentally route requests through alternative AI providers using claude-code-router. This mode is not officially supported and is intended primarily for: diff --git a/docker-compose.yml b/docker-compose.yml index 591b075..843e752 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -27,6 +27,10 @@ services: - CLAUDE_CODE_USE_BEDROCK=${CLAUDE_CODE_USE_BEDROCK:-} - AWS_REGION=${AWS_REGION:-} - AWS_BEARER_TOKEN_BEDROCK=${AWS_BEARER_TOKEN_BEDROCK:-} + - CLAUDE_CODE_USE_VERTEX=${CLAUDE_CODE_USE_VERTEX:-} + - CLOUD_ML_REGION=${CLOUD_ML_REGION:-} + - ANTHROPIC_VERTEX_PROJECT_ID=${ANTHROPIC_VERTEX_PROJECT_ID:-} + - GOOGLE_APPLICATION_CREDENTIALS=${GOOGLE_APPLICATION_CREDENTIALS:-} - ANTHROPIC_SMALL_MODEL=${ANTHROPIC_SMALL_MODEL:-} - ANTHROPIC_MEDIUM_MODEL=${ANTHROPIC_MEDIUM_MODEL:-} - ANTHROPIC_LARGE_MODEL=${ANTHROPIC_LARGE_MODEL:-} @@ -39,6 +43,7 @@ services: - ./prompts:/app/prompts - ./audit-logs:/app/audit-logs - ${OUTPUT_DIR:-./audit-logs}:/app/output + - ./credentials:/app/credentials:ro - ./repos:/repos - ${BENCHMARKS_BASE:-.}:/benchmarks shm_size: 2gb diff --git a/shannon b/shannon index 8606d3b..0a96eba 100755 --- a/shannon +++ b/shannon @@ -156,12 +156,37 @@ cmd_start() { echo "ERROR: Bedrock mode requires the following env vars in .env:$MISSING" exit 1 fi + elif [ "$CLAUDE_CODE_USE_VERTEX" = "1" ]; then + # Vertex AI mode — validate required GCP credentials + MISSING="" + [ -z "$CLOUD_ML_REGION" ] && MISSING="$MISSING CLOUD_ML_REGION" + [ -z "$ANTHROPIC_VERTEX_PROJECT_ID" ] && MISSING="$MISSING ANTHROPIC_VERTEX_PROJECT_ID" + [ -z "$ANTHROPIC_SMALL_MODEL" ] && MISSING="$MISSING ANTHROPIC_SMALL_MODEL" + [ -z "$ANTHROPIC_MEDIUM_MODEL" ] && MISSING="$MISSING ANTHROPIC_MEDIUM_MODEL" + [ -z "$ANTHROPIC_LARGE_MODEL" ] && MISSING="$MISSING ANTHROPIC_LARGE_MODEL" + if [ -n "$MISSING" ]; then + echo "ERROR: Vertex AI mode requires the following env vars in .env:$MISSING" + exit 1 + fi + # Validate service account key file (must be inside ./credentials/ for Docker mount) + if [ -z "$GOOGLE_APPLICATION_CREDENTIALS" ]; then + echo "ERROR: Vertex AI mode requires GOOGLE_APPLICATION_CREDENTIALS in .env" + echo " Place your service account key in ./credentials/ and set:" + echo " GOOGLE_APPLICATION_CREDENTIALS=./credentials/gcp-sa-key.json" + exit 1 + fi + if [ ! -f "$GOOGLE_APPLICATION_CREDENTIALS" ]; then + echo "ERROR: Service account key file not found: $GOOGLE_APPLICATION_CREDENTIALS" + echo " Download a key from the GCP Console (IAM > Service Accounts > Keys)" + exit 1 + fi elif [ "$ROUTER" = "true" ] && { [ -n "$OPENAI_API_KEY" ] || [ -n "$OPENROUTER_API_KEY" ]; }; then # Router mode with alternative provider - set a placeholder for SDK init export ANTHROPIC_API_KEY="router-mode" else echo "ERROR: Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env" echo " (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock," + echo " CLAUDE_CODE_USE_VERTEX=1 for Google Vertex AI," echo " or ROUTER=true with OPENAI_API_KEY or OPENROUTER_API_KEY)" exit 1 fi @@ -222,7 +247,7 @@ cmd_start() { fi # Ensure audit-logs directory exists with write permissions for container user (UID 1001) - mkdir -p ./audit-logs + mkdir -p ./audit-logs ./credentials chmod 777 ./audit-logs # Ensure repo deliverables directory is writable by container user (UID 1001) diff --git a/src/ai/claude-executor.ts b/src/ai/claude-executor.ts index 6a3af0c..04c990c 100644 --- a/src/ai/claude-executor.ts +++ b/src/ai/claude-executor.ts @@ -235,6 +235,10 @@ export async function runClaudePrompt( 'CLAUDE_CODE_USE_BEDROCK', 'AWS_REGION', 'AWS_BEARER_TOKEN_BEDROCK', + 'CLAUDE_CODE_USE_VERTEX', + 'CLOUD_ML_REGION', + 'ANTHROPIC_VERTEX_PROJECT_ID', + 'GOOGLE_APPLICATION_CREDENTIALS', 'ANTHROPIC_SMALL_MODEL', 'ANTHROPIC_MEDIUM_MODEL', 'ANTHROPIC_LARGE_MODEL', diff --git a/src/services/preflight.ts b/src/services/preflight.ts index dd2dc15..d050d05 100644 --- a/src/services/preflight.ts +++ b/src/services/preflight.ts @@ -14,7 +14,7 @@ * Checks run sequentially, cheapest first: * 1. Repository path exists and contains .git * 2. Config file parses and validates (if provided) - * 3. Credentials validate via Claude Agent SDK query (API key, OAuth, or router mode) + * 3. Credentials validate via Claude Agent SDK query (API key, OAuth, Bedrock, Vertex AI, or router mode) */ import fs from 'fs/promises'; @@ -185,11 +185,56 @@ async function validateCredentials( return ok(undefined); } - // 3. Check that at least one credential is present + // 3. Vertex AI mode — validate required GCP credentials are present + if (process.env.CLAUDE_CODE_USE_VERTEX === '1') { + const required = ['CLOUD_ML_REGION', 'ANTHROPIC_VERTEX_PROJECT_ID', 'ANTHROPIC_SMALL_MODEL', 'ANTHROPIC_MEDIUM_MODEL', 'ANTHROPIC_LARGE_MODEL']; + const missing = required.filter(v => !process.env[v]); + if (missing.length > 0) { + return err( + new PentestError( + `Vertex AI mode requires the following env vars in .env: ${missing.join(', ')}`, + 'config', + false, + { missing }, + ErrorCode.AUTH_FAILED + ) + ); + } + // Validate service account credentials file is accessible + const credPath = process.env.GOOGLE_APPLICATION_CREDENTIALS; + if (!credPath) { + return err( + new PentestError( + 'Vertex AI mode requires GOOGLE_APPLICATION_CREDENTIALS pointing to a service account key JSON file', + 'config', + false, + {}, + ErrorCode.AUTH_FAILED + ) + ); + } + try { + await fs.access(credPath); + } catch { + return err( + new PentestError( + `Service account key file not found at: ${credPath}`, + 'config', + false, + { credPath }, + ErrorCode.AUTH_FAILED + ) + ); + } + logger.info('Vertex AI credentials OK'); + return ok(undefined); + } + + // 4. Check that at least one credential is present if (!process.env.ANTHROPIC_API_KEY && !process.env.CLAUDE_CODE_OAUTH_TOKEN) { return err( new PentestError( - 'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock)', + 'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock, or CLAUDE_CODE_USE_VERTEX=1 for Google Vertex AI)', 'config', false, {}, @@ -198,7 +243,7 @@ async function validateCredentials( ); } - // 4. Validate via SDK query + // 5. Validate via SDK query const authType = process.env.CLAUDE_CODE_OAUTH_TOKEN ? 'OAuth token' : 'API key'; logger.info(`Validating ${authType} via SDK...`);