diff --git a/shannon.mjs b/shannon.mjs index 1dc54ee..676b54e 100755 --- a/shannon.mjs +++ b/shannon.mjs @@ -12,7 +12,7 @@ import { createSession, updateSession, getSession, AGENTS } from './src/session- import { runPhase, getGitCommitHash } from './src/checkpoint-manager.js'; // Setup and Deliverables -import { setupLocalRepo, cleanupMCP } from './src/setup/environment.js'; +import { setupLocalRepo } from './src/setup/environment.js'; // AI and Prompts import { runClaudePromptWithRetry } from './src/ai/claude-executor.js'; @@ -23,8 +23,8 @@ import { executePreReconPhase } from './src/phases/pre-recon.js'; import { assembleFinalReport } from './src/phases/reporting.js'; // Utils -import { timingResults, costResults, displayTimingSummary, Timer, formatDuration } from './src/utils/metrics.js'; -import { setupLogging } from './src/utils/logger.js'; +import { timingResults, costResults, displayTimingSummary, Timer } from './src/utils/metrics.js'; +import { formatDuration } from './src/audit/utils.js'; // CLI import { handleDeveloperCommand } from './src/cli/command-handler.js'; @@ -44,21 +44,16 @@ import { // Configure zx to disable timeouts (let tools run as long as needed) $.timeout = 0; -// Global cleanup function for logging -let cleanupLogging = null; - // Setup graceful cleanup on process signals process.on('SIGINT', async () => { console.log(chalk.yellow('\n⚠️ Received SIGINT, cleaning up...')); - await cleanupMCP(); - if (cleanupLogging) await cleanupLogging(); + process.exit(0); }); process.on('SIGTERM', async () => { console.log(chalk.yellow('\n⚠️ Received SIGTERM, cleaning up...')); - await cleanupMCP(); - if (cleanupLogging) await cleanupLogging(); + process.exit(0); }); @@ -136,7 +131,6 @@ async function main(webUrl, repoPath, configPath = null, pipelineTestingMode = f console.log(chalk.gray('Use developer commands to run individual agents:')); console.log(chalk.gray(' ./shannon.mjs --run-agent pre-recon')); console.log(chalk.gray(' ./shannon.mjs --status')); - await cleanupMCP(); process.exit(0); } @@ -182,7 +176,6 @@ async function main(webUrl, repoPath, configPath = null, pipelineTestingMode = f if (!nextAgent) { console.log(chalk.green(`✅ All agents completed! Session is finished.`)); await displayTimingSummary(timingResults, costResults, session.completedAgents); - await cleanupMCP(); process.exit(0); } @@ -360,7 +353,6 @@ if (args[0] && args[0].includes('shannon.mjs')) { // Parse flags and arguments let configPath = null; let pipelineTestingMode = false; -let logFilePath = null; const nonFlagArgs = []; let developerCommand = null; const developerCommands = ['--run-phase', '--run-all', '--rollback-to', '--rerun', '--status', '--list-agents', '--cleanup']; @@ -374,16 +366,6 @@ for (let i = 0; i < args.length; i++) { console.log(chalk.red('❌ --config flag requires a file path')); process.exit(1); } - } else if (args[i] === '--log') { - // --log can optionally take a file path, otherwise use default - if (i + 1 < args.length && !args[i + 1].startsWith('-')) { - logFilePath = args[i + 1]; - i++; // Skip the next argument - } else { - // Generate default log filename with timestamp - const timestamp = new Date().toISOString().replace(/[:.]/g, '-').slice(0, -5); - logFilePath = `shannon-${timestamp}.log`; - } } else if (args[i] === '--pipeline-testing') { pipelineTestingMode = true; } else if (developerCommands.includes(args[i])) { @@ -410,25 +392,10 @@ if (args.includes('--help') || args.includes('-h') || args.includes('help')) { process.exit(0); } -// Setup logging if --log flag is present -if (logFilePath) { - try { - cleanupLogging = await setupLogging(logFilePath); - const absoluteLogPath = path.isAbsolute(logFilePath) - ? logFilePath - : path.join(process.cwd(), logFilePath); - console.log(chalk.green(`📝 Logging enabled: ${absoluteLogPath}`)); - } catch (error) { - console.log(chalk.yellow(`⚠️ Failed to setup logging: ${error.message}`)); - console.log(chalk.gray('Continuing without logging...')); - } -} - // Handle developer commands if (developerCommand) { await handleDeveloperCommand(developerCommand, nonFlagArgs, pipelineTestingMode, runClaudePromptWithRetry, loadPrompt); - await cleanupMCP(); - if (cleanupLogging) await cleanupLogging(); + process.exit(0); } @@ -478,8 +445,7 @@ try { const finalReportPath = await main(webUrl, repoPathValidation.path, configPath, pipelineTestingMode); console.log(chalk.green.bold('\n📄 FINAL REPORT AVAILABLE:')); console.log(chalk.cyan(finalReportPath)); - await cleanupMCP(); - if (cleanupLogging) await cleanupLogging(); + } catch (error) { // Enhanced error boundary with proper logging if (error instanceof PentestError) { @@ -499,7 +465,6 @@ try { console.log(chalk.gray(` Stack: ${error?.stack || 'No stack trace available'}`)); } } - await cleanupMCP(); - if (cleanupLogging) await cleanupLogging(); + process.exit(1); } \ No newline at end of file diff --git a/src/ai/claude-executor.js b/src/ai/claude-executor.js index 947ee30..d494cbe 100644 --- a/src/ai/claude-executor.js +++ b/src/ai/claude-executor.js @@ -6,7 +6,8 @@ import { dirname } from 'path'; import { isRetryableError, getRetryDelay, PentestError } from '../error-handling.js'; import { ProgressIndicator } from '../progress-indicator.js'; -import { timingResults, costResults, Timer, formatDuration } from '../utils/metrics.js'; +import { timingResults, costResults, Timer } from '../utils/metrics.js'; +import { formatDuration } from '../audit/utils.js'; import { createGitCheckpoint, commitGitSuccess, rollbackGitWorkspace } from '../utils/git-manager.js'; import { AGENT_VALIDATORS, MCP_AGENT_MAPPING } from '../constants.js'; import { filterJsonToolCalls, getAgentPrefix } from '../utils/output-formatter.js'; diff --git a/src/audit/audit-session.js b/src/audit/audit-session.js index d53da92..c8df674 100644 --- a/src/audit/audit-session.js +++ b/src/audit/audit-session.js @@ -8,32 +8,7 @@ import { AgentLogger } from './logger.js'; import { MetricsTracker } from './metrics-tracker.js'; import { initializeAuditStructure, formatTimestamp } from './utils.js'; - -/** - * SessionMutex for concurrency control - * (Identical to session-manager.js implementation) - */ -class SessionMutex { - constructor() { - this.locks = new Map(); - } - - async lock(sessionId) { - if (this.locks.has(sessionId)) { - // Wait for existing lock to be released - await this.locks.get(sessionId); - } - - let resolve; - const promise = new Promise(r => resolve = r); - this.locks.set(sessionId, promise); - - return () => { - this.locks.delete(sessionId); - resolve(); - }; - } -} +import { SessionMutex } from '../utils/concurrency.js'; // Global mutex instance const sessionMutex = new SessionMutex(); @@ -100,31 +75,6 @@ export class AuditSession { } } - /** - * Log session-level failure (pre-agent failures) - * @param {Error} error - Error object - * @param {Object} context - Additional context - * @returns {Promise} - */ - async logSessionFailure(error, context = {}) { - await this.ensureInitialized(); - - // Update session status - await this.metricsTracker.updateSessionStatus('failed'); - - // Create a special failure logger - const failureLogger = new AgentLogger(this.sessionMetadata, 'session-failure', 1); - await failureLogger.initialize(); - - await failureLogger.logError(error, { - ...context, - timestamp: formatTimestamp(), - sessionId: this.sessionId - }); - - await failureLogger.close(); - } - /** * Start agent execution * @param {string} agentName - Agent name @@ -169,19 +119,6 @@ export class AuditSession { await this.currentLogger.logEvent(eventType, eventData); } - /** - * Log a text message (for compatibility) - * @param {string} message - Message to log - * @returns {Promise} - */ - async logMessage(message) { - if (!this.currentLogger) { - throw new Error('No active logger. Call startAgent() first.'); - } - - await this.currentLogger.logMessage(message); - } - /** * End agent execution (mutex-protected) * @param {string} agentName - Agent name @@ -224,41 +161,6 @@ export class AuditSession { } } - /** - * Update validation results - * @param {string} agentName - Agent name - * @param {Object} validationData - Validation data - * @returns {Promise} - */ - async updateValidation(agentName, validationData) { - await this.ensureInitialized(); - - const unlock = await sessionMutex.lock(this.sessionId); - try { - await this.metricsTracker.reload(); - await this.metricsTracker.updateValidation(agentName, validationData); - } finally { - unlock(); - } - } - - /** - * Mark agent as rolled back - * @param {string} agentName - Agent name - * @returns {Promise} - */ - async markRolledBack(agentName) { - await this.ensureInitialized(); - - const unlock = await sessionMutex.lock(this.sessionId); - try { - await this.metricsTracker.reload(); - await this.metricsTracker.markRolledBack(agentName); - } finally { - unlock(); - } - } - /** * Mark multiple agents as rolled back * @param {string[]} agentNames - Array of agent names @@ -301,13 +203,4 @@ export class AuditSession { await this.ensureInitialized(); return this.metricsTracker.getMetrics(); } - - /** - * Get validation results (read-only) - * @returns {Promise} Validation results - */ - async getValidation() { - await this.ensureInitialized(); - return this.metricsTracker.getValidation(); - } } diff --git a/src/audit/logger.js b/src/audit/logger.js index 092d705..bb03ab2 100644 --- a/src/audit/logger.js +++ b/src/audit/logger.js @@ -124,81 +124,6 @@ export class AgentLogger { return this.writeRaw(eventLine); } - /** - * Log a text message (for compatibility with existing logging) - * @param {string} message - Message to log - * @returns {Promise} - */ - async logMessage(message) { - const timestamp = formatTimestamp(); - const line = `[${timestamp}] ${message}\n`; - return this.writeRaw(line); - } - - /** - * Log tool start event - * @param {string} toolName - Name of the tool - * @param {Object} [parameters] - Tool parameters - * @returns {Promise} - */ - async logToolStart(toolName, parameters = {}) { - return this.logEvent('tool_start', { toolName, parameters }); - } - - /** - * Log tool end event - * @param {string} toolName - Name of the tool - * @param {Object} result - Tool result - * @returns {Promise} - */ - async logToolEnd(toolName, result) { - return this.logEvent('tool_end', { toolName, result }); - } - - /** - * Log LLM response event - * @param {string} content - Response content - * @param {Object} [metadata] - Additional metadata - * @returns {Promise} - */ - async logLLMResponse(content, metadata = {}) { - return this.logEvent('llm_response', { content, ...metadata }); - } - - /** - * Log validation start event - * @param {string} validationType - Type of validation - * @returns {Promise} - */ - async logValidationStart(validationType) { - return this.logEvent('validation_start', { validationType }); - } - - /** - * Log validation end event - * @param {string} validationType - Type of validation - * @param {boolean} success - Whether validation passed - * @param {Object} [details] - Validation details - * @returns {Promise} - */ - async logValidationEnd(validationType, success, details = {}) { - return this.logEvent('validation_end', { validationType, success, ...details }); - } - - /** - * Log error event - * @param {Error} error - Error object - * @param {Object} [context] - Additional context - * @returns {Promise} - */ - async logError(error, context = {}) { - return this.logEvent('error', { - message: error.message, - stack: error.stack, - ...context - }); - } - /** * Close the log stream * @returns {Promise} diff --git a/src/audit/utils.js b/src/audit/utils.js index 751eddf..58a2f05 100644 --- a/src/audit/utils.js +++ b/src/audit/utils.js @@ -138,15 +138,6 @@ export function formatDuration(ms) { return `${minutes}m ${remainingSeconds}s`; } -/** - * Format cost in USD - * @param {number} usd - Cost in USD - * @returns {string} Formatted cost (e.g., "$0.0823", "$2.14") - */ -export function formatCost(usd) { - return `$${usd.toFixed(4)}`; -} - /** * Format timestamp to ISO 8601 string * @param {number} [timestamp] - Unix timestamp in ms (defaults to now) diff --git a/src/checkpoint-manager.js b/src/checkpoint-manager.js index 0685ad4..a15be0e 100644 --- a/src/checkpoint-manager.js +++ b/src/checkpoint-manager.js @@ -3,6 +3,7 @@ import chalk from 'chalk'; import { PentestError } from './error-handling.js'; import { parseConfig, distributeConfig } from './config-parser.js'; import { executeGitCommandWithRetry } from './utils/git-manager.js'; +import { formatDuration } from './audit/utils.js'; import { AGENTS, PHASES, @@ -900,23 +901,3 @@ const getTimeAgo = (timestamp) => { } }; -// Helper function to format duration in milliseconds to human readable format -const formatDuration = (durationMs) => { - if (durationMs < 1000) { - return `${durationMs}ms`; - } - - const seconds = Math.floor(durationMs / 1000); - const minutes = Math.floor(seconds / 60); - const hours = Math.floor(minutes / 60); - - if (hours > 0) { - return `${hours}h ${minutes % 60}m ${seconds % 60}s`; - } else if (minutes > 0) { - return `${minutes}m ${seconds % 60}s`; - } else { - return `${seconds}s`; - } -}; - - diff --git a/src/cli/command-handler.js b/src/cli/command-handler.js index 47b9cf0..08629c5 100644 --- a/src/cli/command-handler.js +++ b/src/cli/command-handler.js @@ -7,7 +7,7 @@ import { runPhase, runAll, rollbackTo, rerunAgent, displayStatus, listAgents } from '../checkpoint-manager.js'; import { logError, PentestError } from '../error-handling.js'; -import { cleanupMCP } from '../setup/environment.js'; +import { promptConfirmation } from './prompts.js'; // Developer command handlers export async function handleDeveloperCommand(command, args, pipelineTestingMode, runClaudePromptWithRetry, loadPrompt) { @@ -27,41 +27,19 @@ export async function handleDeveloperCommand(command, args, pipelineTestingMode, const sessionId = args[0]; const deletedSession = await deleteSession(sessionId); console.log(chalk.green(`✅ Deleted session ${sessionId} (${new URL(deletedSession.webUrl).hostname})`)); - // Clean up MCP agents when deleting specific session - await cleanupMCP(); } else { // Cleanup all sessions - require confirmation - console.log(chalk.yellow('⚠️ This will delete all pentest sessions. Are you sure? (y/N):')); - const { createInterface } = await import('readline'); - const readline = createInterface({ - input: process.stdin, - output: process.stdout - }); - - await new Promise((resolve) => { - readline.question('', (answer) => { - readline.close(); - if (answer.toLowerCase() === 'y' || answer.toLowerCase() === 'yes') { - deleteAllSessions().then(deleted => { - if (deleted) { - console.log(chalk.green('✅ All sessions deleted')); - } else { - console.log(chalk.yellow('⚠️ No sessions found to delete')); - } - // Clean up MCP agents after deleting sessions - return cleanupMCP(); - }).then(() => { - resolve(); - }).catch(error => { - console.log(chalk.red(`❌ Failed to delete sessions: ${error.message}`)); - resolve(); - }); - } else { - console.log(chalk.gray('Cleanup cancelled')); - resolve(); - } - }); - }); + const confirmed = await promptConfirmation(chalk.yellow('⚠️ This will delete all pentest sessions. Are you sure? (y/N):')); + if (confirmed) { + const deleted = await deleteAllSessions(); + if (deleted) { + console.log(chalk.green('✅ All sessions deleted')); + } else { + console.log(chalk.yellow('⚠️ No sessions found to delete')); + } + } else { + console.log(chalk.gray('Cleanup cancelled')); + } } return; } diff --git a/src/cli/prompts.js b/src/cli/prompts.js new file mode 100644 index 0000000..5d174e8 --- /dev/null +++ b/src/cli/prompts.js @@ -0,0 +1,62 @@ +import { createInterface } from 'readline'; +import { PentestError } from '../error-handling.js'; + +/** + * Prompt user for yes/no confirmation + * @param {string} message - Question to display + * @returns {Promise} true if confirmed, false otherwise + */ +export async function promptConfirmation(message) { + const readline = createInterface({ + input: process.stdin, + output: process.stdout + }); + + return new Promise((resolve) => { + readline.question(message + ' ', (answer) => { + readline.close(); + const confirmed = answer.toLowerCase() === 'y' || answer.toLowerCase() === 'yes'; + resolve(confirmed); + }); + }); +} + +/** + * Prompt user to select from numbered list + * @param {string} message - Selection prompt + * @param {Array} items - Items to choose from + * @returns {Promise} Selected item + * @throws {PentestError} If invalid selection + */ +export async function promptSelection(message, items) { + if (!items || items.length === 0) { + throw new PentestError( + 'No items available for selection', + 'validation', + false + ); + } + + const readline = createInterface({ + input: process.stdin, + output: process.stdout + }); + + return new Promise((resolve, reject) => { + readline.question(message + ' ', (answer) => { + readline.close(); + + const choice = parseInt(answer); + if (isNaN(choice) || choice < 1 || choice > items.length) { + reject(new PentestError( + `Invalid selection. Please enter a number between 1 and ${items.length}`, + 'validation', + false, + { choice: answer } + )); + } else { + resolve(items[choice - 1]); + } + }); + }); +} diff --git a/src/cli/ui.js b/src/cli/ui.js index cea967e..c12a4c0 100644 --- a/src/cli/ui.js +++ b/src/cli/ui.js @@ -21,7 +21,6 @@ export function showHelp() { console.log(chalk.yellow.bold('OPTIONS:')); console.log(' --config YAML configuration file for authentication and testing parameters'); - console.log(' --log [file] Capture all output to log file (default: shannon-.log)'); console.log(' --pipeline-testing Use minimal prompts for fast pipeline testing (creates minimal deliverables)\n'); console.log(chalk.yellow.bold('DEVELOPER COMMANDS:')); @@ -37,7 +36,6 @@ export function showHelp() { console.log(' # Normal mode - create new session'); console.log(' ./shannon.mjs "https://example.com" "/path/to/local/repo"'); console.log(' ./shannon.mjs "https://example.com" "/path/to/local/repo" --config auth.yaml'); - console.log(' ./shannon.mjs "https://example.com" "/path/to/local/repo" --log pentest.log'); console.log(' ./shannon.mjs "https://example.com" "/path/to/local/repo" --setup-only # Setup only\n'); console.log(' # Developer mode - operate on existing session'); diff --git a/src/constants.js b/src/constants.js index 5e7c16f..40cc994 100644 --- a/src/constants.js +++ b/src/constants.js @@ -2,6 +2,27 @@ import { path, fs } from 'zx'; import chalk from 'chalk'; import { validateQueueAndDeliverable } from './queue-validation.js'; +// Factory function for vulnerability queue validators +function createVulnValidator(vulnType) { + return async (sourceDir) => { + try { + await validateQueueAndDeliverable(vulnType, sourceDir); + return true; + } catch (error) { + console.log(chalk.yellow(` Queue validation failed for ${vulnType}: ${error.message}`)); + return false; + } + }; +} + +// Factory function for exploit deliverable validators +function createExploitValidator(vulnType) { + return async (sourceDir) => { + const evidenceFile = path.join(sourceDir, 'deliverables', `${vulnType}_exploitation_evidence.md`); + return await fs.pathExists(evidenceFile); + }; +} + // MCP agent mapping - assigns each agent to a specific Playwright instance to prevent conflicts export const MCP_AGENT_MAPPING = Object.freeze({ // Phase 1: Pre-reconnaissance (actual prompt name is 'pre-recon-code') @@ -47,81 +68,18 @@ export const AGENT_VALIDATORS = Object.freeze({ }, // Vulnerability analysis agents - 'injection-vuln': async (sourceDir) => { - try { - await validateQueueAndDeliverable('injection', sourceDir); - return true; - } catch (error) { - console.log(chalk.yellow(` Queue validation failed for injection: ${error.message}`)); - return false; - } - }, - - 'xss-vuln': async (sourceDir) => { - try { - await validateQueueAndDeliverable('xss', sourceDir); - return true; - } catch (error) { - console.log(chalk.yellow(` Queue validation failed for xss: ${error.message}`)); - return false; - } - }, - - 'auth-vuln': async (sourceDir) => { - try { - await validateQueueAndDeliverable('auth', sourceDir); - return true; - } catch (error) { - console.log(chalk.yellow(` Queue validation failed for auth: ${error.message}`)); - return false; - } - }, - - 'ssrf-vuln': async (sourceDir) => { - try { - await validateQueueAndDeliverable('ssrf', sourceDir); - return true; - } catch (error) { - console.log(chalk.yellow(` Queue validation failed for ssrf: ${error.message}`)); - return false; - } - }, - - 'authz-vuln': async (sourceDir) => { - try { - await validateQueueAndDeliverable('authz', sourceDir); - return true; - } catch (error) { - console.log(chalk.yellow(` Queue validation failed for authz: ${error.message}`)); - return false; - } - }, + 'injection-vuln': createVulnValidator('injection'), + 'xss-vuln': createVulnValidator('xss'), + 'auth-vuln': createVulnValidator('auth'), + 'ssrf-vuln': createVulnValidator('ssrf'), + 'authz-vuln': createVulnValidator('authz'), // Exploitation agents - 'injection-exploit': async (sourceDir) => { - const evidenceFile = path.join(sourceDir, 'deliverables', 'injection_exploitation_evidence.md'); - return await fs.pathExists(evidenceFile); - }, - - 'xss-exploit': async (sourceDir) => { - const evidenceFile = path.join(sourceDir, 'deliverables', 'xss_exploitation_evidence.md'); - return await fs.pathExists(evidenceFile); - }, - - 'auth-exploit': async (sourceDir) => { - const evidenceFile = path.join(sourceDir, 'deliverables', 'auth_exploitation_evidence.md'); - return await fs.pathExists(evidenceFile); - }, - - 'ssrf-exploit': async (sourceDir) => { - const evidenceFile = path.join(sourceDir, 'deliverables', 'ssrf_exploitation_evidence.md'); - return await fs.pathExists(evidenceFile); - }, - - 'authz-exploit': async (sourceDir) => { - const evidenceFile = path.join(sourceDir, 'deliverables', 'authz_exploitation_evidence.md'); - return await fs.pathExists(evidenceFile); - }, + 'injection-exploit': createExploitValidator('injection'), + 'xss-exploit': createExploitValidator('xss'), + 'auth-exploit': createExploitValidator('auth'), + 'ssrf-exploit': createExploitValidator('ssrf'), + 'authz-exploit': createExploitValidator('authz'), // Executive report agent 'report': async (sourceDir) => { diff --git a/src/error-handling.js b/src/error-handling.js index dabeae6..5b7f554 100644 --- a/src/error-handling.js +++ b/src/error-handling.js @@ -51,18 +51,6 @@ export const logError = async (error, contextMsg, sourceDir = null) => { return logEntry; }; -// Handle configuration parsing errors -const handleConfigError = (error, configPath) => { - const configError = new PentestError( - `Configuration error in ${configPath}: ${error.message}. Check your config.yaml file format and try again.`, - 'config', - false, - { configPath, originalError: error.message } - ); - throw configError; -}; - - // Handle tool execution errors export const handleToolError = (toolName, error) => { const isRetryable = error.code === 'ECONNRESET' || error.code === 'ETIMEDOUT' || error.code === 'ENOTFOUND'; @@ -167,22 +155,4 @@ export const getRetryDelay = (error, attempt) => { const baseDelay = Math.pow(2, attempt) * 1000; // 2s, 4s, 8s const jitter = Math.random() * 1000; // 0-1s random return Math.min(baseDelay + jitter, 30000); // Max 30s -}; - -// General error handler with context -const handleError = (error, context, isFatal = false) => { - const pentestError = error instanceof PentestError - ? error - : new PentestError(error.message, 'unknown', false, { context, originalError: error.message }); - - if (isFatal) { - pentestError.type = 'fatal'; - throw pentestError; - } - - return { - success: false, - error: pentestError, - continuable: !isFatal - }; }; \ No newline at end of file diff --git a/src/phases/pre-recon.js b/src/phases/pre-recon.js index a7d4c67..9020eb1 100644 --- a/src/phases/pre-recon.js +++ b/src/phases/pre-recon.js @@ -1,6 +1,7 @@ import { $, fs, path } from 'zx'; import chalk from 'chalk'; -import { Timer, timingResults, formatDuration } from '../utils/metrics.js'; +import { Timer, timingResults } from '../utils/metrics.js'; +import { formatDuration } from '../audit/utils.js'; import { handleToolError, PentestError } from '../error-handling.js'; import { AGENTS } from '../session-manager.js'; import { runClaudePromptWithRetry } from '../ai/claude-executor.js'; diff --git a/src/progress-indicator.js b/src/progress-indicator.js index c917b6d..91d1f2b 100644 --- a/src/progress-indicator.js +++ b/src/progress-indicator.js @@ -22,10 +22,6 @@ export class ProgressIndicator { }, 100); } - updateMessage(newMessage) { - this.message = newMessage; - } - stop() { if (!this.isRunning) return; diff --git a/src/queue-validation.js b/src/queue-validation.js index e94d485..819ac2a 100644 --- a/src/queue-validation.js +++ b/src/queue-validation.js @@ -27,7 +27,6 @@ const VULN_TYPE_CONFIG = Object.freeze({ // Functional composition utilities - async pipe for promise chain const pipe = (...fns) => x => fns.reduce(async (v, f) => f(await v), x); -const compose = (...fns) => x => fns.reduceRight((v, f) => f(v), x); // Pure function to create validation rule const createValidationRule = (predicate, errorMessage, retryable = true) => diff --git a/src/session-manager.js b/src/session-manager.js index dfbdfd1..17122c3 100644 --- a/src/session-manager.js +++ b/src/session-manager.js @@ -2,6 +2,8 @@ import { fs, path } from 'zx'; import chalk from 'chalk'; import crypto from 'crypto'; import { PentestError } from './error-handling.js'; +import { SessionMutex } from './utils/concurrency.js'; +import { promptSelection } from './cli/prompts.js'; // Generate a session-based log folder path // NEW FORMAT: {hostname}_{sessionId} (no hash, full UUID for consistency with audit system) @@ -11,29 +13,6 @@ export const generateSessionLogPath = (webUrl, sessionId) => { return path.join(process.cwd(), 'agent-logs', sessionFolderName); }; -// Mutex for session file operations to prevent race conditions -class SessionMutex { - constructor() { - this.locks = new Map(); - } - - async lock(sessionId) { - if (this.locks.has(sessionId)) { - // Wait for existing lock to be released - await this.locks.get(sessionId); - } - - let resolve; - const promise = new Promise(r => resolve = r); - this.locks.set(sessionId, promise); - - return () => { - this.locks.delete(sessionId); - resolve(); - }; - } -} - const sessionMutex = new SessionMutex(); // Agent definitions according to PRD @@ -338,29 +317,10 @@ export const selectSession = async () => { }); // Get user selection - const { createInterface } = await import('readline'); - const readline = createInterface({ - input: process.stdin, - output: process.stdout - }); - - return new Promise((resolve, reject) => { - readline.question(chalk.cyan(`Select session (1-${sessions.length}): `), (answer) => { - readline.close(); - - const choice = parseInt(answer); - if (isNaN(choice) || choice < 1 || choice > sessions.length) { - reject(new PentestError( - `Invalid selection. Please enter a number between 1 and ${sessions.length}`, - 'validation', - false, - { choice: answer } - )); - } else { - resolve(sessions[choice - 1]); - } - }); - }); + return await promptSelection( + chalk.cyan(`Select session (1-${sessions.length}):`), + sessions + ); }; // Validate agent name diff --git a/src/setup/environment.js b/src/setup/environment.js index ccd90a5..5fd8921 100644 --- a/src/setup/environment.js +++ b/src/setup/environment.js @@ -1,91 +1,14 @@ import { $, fs, path } from 'zx'; import chalk from 'chalk'; -import { PentestError, logError } from '../error-handling.js'; - -// Pure function: Setup MCP with multiple isolated Playwright instances -export async function setupMCP(sourceDir) { - console.log(chalk.blue('🎭 Setting up 5 isolated Playwright MCP instances...')); - - // Set headless mode for all instances - process.env.PLAYWRIGHT_HEADLESS = 'true'; - - try { - // Clean slate - remove any existing instances - const instancesToRemove = ['playwright', ...Array.from({length: 5}, (_, i) => `playwright-agent${i + 1}`)]; - - for (const instance of instancesToRemove) { - try { - await $`claude mcp remove ${instance} --scope user 2>/dev/null`; - } catch { - // Silent ignore - instance might not exist - } - } - - // Create 5 isolated instances sequentially to avoid config conflicts - for (let i = 1; i <= 5; i++) { - const instanceName = `playwright-agent${i}`; - const userDataDir = `/tmp/${instanceName}`; - - // Ensure user data directory exists - await fs.ensureDir(userDataDir); - - try { - await $`claude mcp add ${instanceName} --scope user -- npx @playwright/mcp@latest --isolated --user-data-dir ${userDataDir}`; - console.log(chalk.green(` ✅ ${instanceName} configured`)); - } catch (error) { - if (error.message?.includes('already exists')) { - console.log(chalk.gray(` ⏭️ ${instanceName} already exists`)); - } else { - console.log(chalk.yellow(` ⚠️ ${instanceName} failed: ${error.message}, continuing...`)); - } - } - } - console.log(chalk.green('✅ All 5 Playwright MCP instances ready for parallel execution')); - - } catch (error) { - // All MCP setup failures are fatal - const mcpError = new PentestError( - `Critical MCP setup failure: ${error.message}. Browser automation required for pentesting.`, - 'tool', - false, - { sourceDir, originalError: error.message } - ); - await logError(mcpError, 'MCP setup failure', sourceDir); - throw mcpError; - } -} - -// Pure function: Cleanup MCP instances -export async function cleanupMCP() { - console.log(chalk.blue('🧹 Cleaning up Playwright MCP instances...')); - - try { - // Remove all instances (including legacy 'playwright' if it exists) - const instancesToRemove = ['playwright', ...Array.from({length: 5}, (_, i) => `playwright-agent${i + 1}`)]; - - for (const instance of instancesToRemove) { - try { - await $`claude mcp remove ${instance} --scope user 2>/dev/null`; - console.log(chalk.gray(` 🗑️ Removed ${instance}`)); - } catch { - // Silent ignore - instance might not exist - } - } - console.log(chalk.green('✅ Playwright MCP cleanup complete')); - - } catch (error) { - // Non-fatal - log warning but don't throw - console.log(chalk.yellow(`⚠️ MCP cleanup warning: ${error.message}`)); - } -} +import { PentestError } from '../error-handling.js'; // Pure function: Setup local repository for testing export async function setupLocalRepo(repoPath) { try { const sourceDir = path.resolve(repoPath); - // Setup MCP in the local repository - critical for browser automation - await setupMCP(sourceDir); + // MCP servers are now configured via mcpServers option in claude-executor.js + // No need for pre-setup with claude CLI // Initialize git repository if not already initialized and create checkpoint try { diff --git a/src/tool-checker.js b/src/tool-checker.js index 30785a8..f6a7f2d 100644 --- a/src/tool-checker.js +++ b/src/tool-checker.js @@ -48,17 +48,6 @@ export const handleMissingTools = (toolAvailability) => { }); console.log(''); } - + return missing; }; - -// Check if a specific tool is available -const isToolAvailable = async (toolName) => { - try { - await $`command -v ${toolName}`; - return true; - } catch { - return false; - } -}; - diff --git a/src/utils/concurrency.js b/src/utils/concurrency.js new file mode 100644 index 0000000..b29b644 --- /dev/null +++ b/src/utils/concurrency.js @@ -0,0 +1,54 @@ +/** + * Concurrency Control Utilities + * + * Provides mutex implementation for preventing race conditions during + * concurrent session operations. + */ + +/** + * SessionMutex - Promise-based mutex for session file operations + * + * Prevents race conditions when multiple agents or operations attempt to + * modify the same session data simultaneously. This is particularly important + * during parallel execution of vulnerability analysis and exploitation phases. + * + * Usage: + * ```js + * const mutex = new SessionMutex(); + * const unlock = await mutex.lock(sessionId); + * try { + * // Critical section - modify session data + * } finally { + * unlock(); // Always release the lock + * } + * ``` + */ +export class SessionMutex { + constructor() { + // Map of sessionId -> Promise (represents active lock) + this.locks = new Map(); + } + + /** + * Acquire lock for a session + * @param {string} sessionId - Session ID to lock + * @returns {Promise} Unlock function to release the lock + */ + async lock(sessionId) { + if (this.locks.has(sessionId)) { + // Wait for existing lock to be released + await this.locks.get(sessionId); + } + + // Create new lock promise + let resolve; + const promise = new Promise(r => resolve = r); + this.locks.set(sessionId, promise); + + // Return unlock function + return () => { + this.locks.delete(sessionId); + resolve(); + }; + } +} diff --git a/src/utils/logger.js b/src/utils/logger.js deleted file mode 100644 index be285d4..0000000 --- a/src/utils/logger.js +++ /dev/null @@ -1,126 +0,0 @@ -import { fs } from 'zx'; -import { path } from 'zx'; - -/** - * Strips ANSI escape codes from a string - * @param {string} str - String with ANSI codes - * @returns {string} Clean string without ANSI codes - */ -function stripAnsi(str) { - if (typeof str !== 'string') { - return str; - } - - // Remove ANSI escape sequences - // This regex matches all common ANSI codes including: - // - Colors (e.g., \x1b[32m) - // - Cursor movement (e.g., \x1b[1;1H) - // - Screen clearing (e.g., \x1b[0J) - // - 256-color codes (e.g., \x1b[38;2;244;197;66m) - return str.replace( - // eslint-disable-next-line no-control-regex - /\x1b\[[0-9;]*[a-zA-Z]|\x1b\][0-9];.*?\x07|\x1b\[[\d;]*m/g, - '' - ); -} - -/** - * Sets up logging to capture all stdout and stderr to a file - * @param {string} logFilePath - Path to the log file - * @returns {Promise} Cleanup function to restore original streams - */ -export async function setupLogging(logFilePath) { - // Resolve to absolute path - const absoluteLogPath = path.isAbsolute(logFilePath) - ? logFilePath - : path.join(process.cwd(), logFilePath); - - // Ensure the directory exists - await fs.ensureDir(path.dirname(absoluteLogPath)); - - // Create write stream for the log file - const logStream = fs.createWriteStream(absoluteLogPath, { flags: 'a' }); - - // Buffer for lines that might be overwritten (carriage return without newline) - let stdoutBuffer = ''; - let stderrBuffer = ''; - - // Store original stdout/stderr write functions - const originalStdoutWrite = process.stdout.write.bind(process.stdout); - const originalStderrWrite = process.stderr.write.bind(process.stderr); - - // Override stdout - process.stdout.write = function(chunk, encoding, callback) { - // Write colorized output to terminal - originalStdoutWrite(chunk, encoding, callback); - - // Write plain text (without ANSI codes) to log file - const cleanChunk = stripAnsi(chunk.toString()); - - // Handle carriage returns - only log when we get a newline - if (cleanChunk.includes('\r') && !cleanChunk.includes('\n')) { - // Buffer this line - it will be overwritten in terminal - stdoutBuffer = cleanChunk.replace(/\r/g, ''); - } else if (cleanChunk.includes('\n')) { - // Flush buffer if exists, then write the new line - if (stdoutBuffer) { - stdoutBuffer = ''; // Clear buffer without writing (it was overwritten) - } - logStream.write(cleanChunk); - } else { - // Normal write - logStream.write(cleanChunk); - } - - return true; - }; - - // Override stderr - process.stderr.write = function(chunk, encoding, callback) { - // Write colorized output to terminal - originalStderrWrite(chunk, encoding, callback); - - // Write plain text (without ANSI codes) to log file - const cleanChunk = stripAnsi(chunk.toString()); - - // Handle carriage returns - only log when we get a newline - if (cleanChunk.includes('\r') && !cleanChunk.includes('\n')) { - // Buffer this line - it will be overwritten in terminal - stderrBuffer = cleanChunk.replace(/\r/g, ''); - } else if (cleanChunk.includes('\n')) { - // Flush buffer if exists, then write the new line - if (stderrBuffer) { - stderrBuffer = ''; // Clear buffer without writing (it was overwritten) - } - logStream.write(cleanChunk); - } else { - // Normal write - logStream.write(cleanChunk); - } - - return true; - }; - - // Return cleanup function - return async function cleanup() { - // Restore original streams - process.stdout.write = originalStdoutWrite; - process.stderr.write = originalStderrWrite; - - // Flush any remaining buffers - if (stdoutBuffer) { - logStream.write(stdoutBuffer + '\n'); - } - if (stderrBuffer) { - logStream.write(stderrBuffer + '\n'); - } - - // Close the log stream - return new Promise((resolve, reject) => { - logStream.end((err) => { - if (err) reject(err); - else resolve(); - }); - }); - }; -} diff --git a/src/utils/metrics.js b/src/utils/metrics.js index e91879e..dc55865 100644 --- a/src/utils/metrics.js +++ b/src/utils/metrics.js @@ -1,13 +1,7 @@ import chalk from 'chalk'; +import { formatDuration } from '../audit/utils.js'; // Timing utilities -export const formatDuration = (ms) => { - if (ms < 1000) return `${ms}ms`; - if (ms < 60000) return `${(ms / 1000).toFixed(1)}s`; - const minutes = Math.floor(ms / 60000); - const seconds = Math.floor((ms % 60000) / 1000); - return `${minutes}m ${seconds}s`; -}; export class Timer { constructor(name) {