ajmallesh
@@ -0,0 +1,242 @@
|
||||
import { fs, path } from 'zx';
|
||||
import chalk from 'chalk';
|
||||
import { PentestError, handlePromptError } from '../error-handling.js';
|
||||
import { MCP_AGENT_MAPPING } from '../constants.js';
|
||||
|
||||
// Pure function: Build complete login instructions from config
|
||||
async function buildLoginInstructions(authentication) {
|
||||
try {
|
||||
// Load the login instructions template
|
||||
const loginInstructionsPath = path.join(import.meta.dirname, '..', '..', 'login_resources', 'login_instructions.txt');
|
||||
|
||||
if (!await fs.pathExists(loginInstructionsPath)) {
|
||||
throw new PentestError(
|
||||
'Login instructions template not found',
|
||||
'filesystem',
|
||||
false,
|
||||
{ loginInstructionsPath }
|
||||
);
|
||||
}
|
||||
|
||||
const fullTemplate = await fs.readFile(loginInstructionsPath, 'utf8');
|
||||
|
||||
// Helper function to extract sections based on markers
|
||||
const getSection = (content, sectionName) => {
|
||||
const regex = new RegExp(`<!-- BEGIN:${sectionName} -->([\\s\\S]*?)<!-- END:${sectionName} -->`, 'g');
|
||||
const match = regex.exec(content);
|
||||
return match ? match[1].trim() : '';
|
||||
};
|
||||
|
||||
// Extract sections based on login type
|
||||
const loginType = authentication.login_type?.toUpperCase();
|
||||
let loginInstructions = '';
|
||||
|
||||
// Build instructions with only relevant sections
|
||||
const commonSection = getSection(fullTemplate, 'COMMON');
|
||||
const authSection = getSection(fullTemplate, loginType); // FORM or SSO
|
||||
const verificationSection = getSection(fullTemplate, 'VERIFICATION');
|
||||
|
||||
// Fallback to full template if markers are missing (backward compatibility)
|
||||
if (!commonSection && !authSection && !verificationSection) {
|
||||
console.log(chalk.yellow('⚠️ Section markers not found, using full login instructions template'));
|
||||
loginInstructions = fullTemplate;
|
||||
} else {
|
||||
// Combine relevant sections
|
||||
loginInstructions = [commonSection, authSection, verificationSection]
|
||||
.filter(section => section) // Remove empty sections
|
||||
.join('\n\n');
|
||||
}
|
||||
|
||||
// Replace the user instructions placeholder with the login flow from config
|
||||
let userInstructions = authentication.login_flow.join('\n');
|
||||
|
||||
// Replace credential placeholders within the user instructions
|
||||
if (authentication.credentials) {
|
||||
if (authentication.credentials.username) {
|
||||
userInstructions = userInstructions.replace(/\$username/g, authentication.credentials.username);
|
||||
}
|
||||
if (authentication.credentials.password) {
|
||||
userInstructions = userInstructions.replace(/\$password/g, authentication.credentials.password);
|
||||
}
|
||||
if (authentication.credentials.totp_secret) {
|
||||
userInstructions = userInstructions.replace(/\$totp/g, `generated TOTP code using secret "${authentication.credentials.totp_secret}"`);
|
||||
}
|
||||
}
|
||||
|
||||
loginInstructions = loginInstructions.replace(/{{user_instructions}}/g, userInstructions);
|
||||
|
||||
// Replace TOTP secret placeholder if present in template
|
||||
if (authentication.credentials?.totp_secret) {
|
||||
loginInstructions = loginInstructions.replace(/{{totp_secret}}/g, authentication.credentials.totp_secret);
|
||||
}
|
||||
|
||||
return loginInstructions;
|
||||
} catch (error) {
|
||||
if (error instanceof PentestError) {
|
||||
throw error;
|
||||
}
|
||||
throw new PentestError(
|
||||
`Failed to build login instructions: ${error.message}`,
|
||||
'config',
|
||||
false,
|
||||
{ authentication, originalError: error.message }
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Pure function: Variable interpolation
|
||||
async function interpolateVariables(template, variables, config = null) {
|
||||
try {
|
||||
if (!template || typeof template !== 'string') {
|
||||
throw new PentestError(
|
||||
'Template must be a non-empty string',
|
||||
'validation',
|
||||
false,
|
||||
{ templateType: typeof template, templateLength: template?.length }
|
||||
);
|
||||
}
|
||||
|
||||
if (!variables || !variables.webUrl || !variables.repoPath) {
|
||||
throw new PentestError(
|
||||
'Variables must include webUrl and repoPath',
|
||||
'validation',
|
||||
false,
|
||||
{ variables: Object.keys(variables || {}) }
|
||||
);
|
||||
}
|
||||
|
||||
let result = template
|
||||
.replace(/{{WEB_URL}}/g, variables.webUrl)
|
||||
.replace(/{{REPO_PATH}}/g, variables.repoPath)
|
||||
.replace(/{{MCP_SERVER}}/g, variables.MCP_SERVER || 'playwright-agent1');
|
||||
|
||||
if (config) {
|
||||
// Handle rules section - if both are empty, use cleaner messaging
|
||||
const hasAvoidRules = config.avoid && config.avoid.length > 0;
|
||||
const hasFocusRules = config.focus && config.focus.length > 0;
|
||||
|
||||
if (!hasAvoidRules && !hasFocusRules) {
|
||||
// Replace the entire rules section with a clean message
|
||||
const cleanRulesSection = '<rules>\nNo specific rules or focus areas provided for this test.\n</rules>';
|
||||
result = result.replace(/<rules>[\s\S]*?<\/rules>/g, cleanRulesSection);
|
||||
} else {
|
||||
const avoidRules = hasAvoidRules ? config.avoid.map(r => `- ${r.description}`).join('\n') : 'None';
|
||||
const focusRules = hasFocusRules ? config.focus.map(r => `- ${r.description}`).join('\n') : 'None';
|
||||
|
||||
result = result
|
||||
.replace(/{{RULES_AVOID}}/g, avoidRules)
|
||||
.replace(/{{RULES_FOCUS}}/g, focusRules);
|
||||
}
|
||||
|
||||
// Extract and inject login instructions from config
|
||||
if (config.authentication?.login_flow) {
|
||||
const loginInstructions = await buildLoginInstructions(config.authentication);
|
||||
result = result.replace(/{{LOGIN_INSTRUCTIONS}}/g, loginInstructions);
|
||||
} else {
|
||||
result = result.replace(/{{LOGIN_INSTRUCTIONS}}/g, '');
|
||||
}
|
||||
} else {
|
||||
// Replace the entire rules section with a clean message when no config provided
|
||||
const cleanRulesSection = '<rules>\nNo specific rules or focus areas provided for this test.\n</rules>';
|
||||
result = result.replace(/<rules>[\s\S]*?<\/rules>/g, cleanRulesSection);
|
||||
result = result.replace(/{{LOGIN_INSTRUCTIONS}}/g, '');
|
||||
}
|
||||
|
||||
// Validate that all placeholders have been replaced (excluding instructional text)
|
||||
const remainingPlaceholders = result.match(/\{\{[^}]+\}\}/g);
|
||||
if (remainingPlaceholders) {
|
||||
console.log(chalk.yellow(`⚠️ Warning: Found unresolved placeholders in prompt: ${remainingPlaceholders.join(', ')}`));
|
||||
}
|
||||
|
||||
return result;
|
||||
} catch (error) {
|
||||
if (error instanceof PentestError) {
|
||||
throw error;
|
||||
}
|
||||
throw new PentestError(
|
||||
`Variable interpolation failed: ${error.message}`,
|
||||
'prompt',
|
||||
false,
|
||||
{ originalError: error.message }
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Pure function: Load and interpolate prompt template
|
||||
export async function loadPrompt(promptName, variables, config = null, pipelineTestingMode = false) {
|
||||
try {
|
||||
// Use pipeline testing prompts if pipeline testing mode is enabled
|
||||
const baseDir = pipelineTestingMode ? 'prompts/pipeline-testing' : 'prompts';
|
||||
const promptsDir = path.join(import.meta.dirname, '..', '..', baseDir);
|
||||
const promptPath = path.join(promptsDir, `${promptName}.txt`);
|
||||
|
||||
// Debug message for pipeline testing mode
|
||||
if (pipelineTestingMode) {
|
||||
console.log(chalk.yellow(`⚡ Using pipeline testing prompt: ${promptPath}`));
|
||||
}
|
||||
|
||||
// Check if file exists first
|
||||
if (!await fs.pathExists(promptPath)) {
|
||||
throw new PentestError(
|
||||
`Prompt file not found: ${promptPath}`,
|
||||
'prompt',
|
||||
false,
|
||||
{ promptName, promptPath }
|
||||
);
|
||||
}
|
||||
|
||||
// Add MCP server assignment to variables
|
||||
const enhancedVariables = { ...variables };
|
||||
|
||||
// Assign MCP server based on prompt name (agent name)
|
||||
if (MCP_AGENT_MAPPING[promptName]) {
|
||||
enhancedVariables.MCP_SERVER = MCP_AGENT_MAPPING[promptName];
|
||||
console.log(chalk.gray(` 🎭 Assigned ${promptName} → ${enhancedVariables.MCP_SERVER}`));
|
||||
} else {
|
||||
// Fallback for unknown agents
|
||||
enhancedVariables.MCP_SERVER = 'playwright-agent1';
|
||||
console.log(chalk.yellow(` 🎭 Unknown agent ${promptName}, using fallback → ${enhancedVariables.MCP_SERVER}`));
|
||||
}
|
||||
|
||||
const template = await fs.readFile(promptPath, 'utf8');
|
||||
return await interpolateVariables(template, enhancedVariables, config);
|
||||
} catch (error) {
|
||||
if (error instanceof PentestError) {
|
||||
throw error;
|
||||
}
|
||||
const promptError = handlePromptError(promptName, error);
|
||||
throw promptError.error;
|
||||
}
|
||||
}
|
||||
|
||||
// Save prompt snapshot for successful agent runs only
|
||||
export async function savePromptSnapshot(sourceDir, agentName, promptContent) {
|
||||
const snapshotDir = path.join(sourceDir, 'prompt-snapshots');
|
||||
await fs.ensureDir(snapshotDir);
|
||||
|
||||
// Use deterministic naming - one snapshot per agent
|
||||
const fileName = `${agentName}.md`;
|
||||
const filePath = path.join(snapshotDir, fileName);
|
||||
|
||||
const timestamp = new Date().toISOString();
|
||||
const snapshotContent = `# Prompt Snapshot: ${agentName}
|
||||
|
||||
**Generated:** ${timestamp}
|
||||
**Agent:** ${agentName}
|
||||
|
||||
---
|
||||
|
||||
## Full Interpolated Prompt
|
||||
|
||||
\`\`\`markdown
|
||||
${promptContent}
|
||||
\`\`\`
|
||||
|
||||
---
|
||||
|
||||
*This snapshot represents the exact prompt that was sent to Claude Code to generate the current deliverables for this agent.*
|
||||
`;
|
||||
|
||||
await fs.writeFile(filePath, snapshotContent);
|
||||
console.log(chalk.gray(` 📸 Prompt snapshot saved: prompt-snapshots/${fileName}`));
|
||||
}
|
||||
Reference in New Issue
Block a user