refactor: consolidate file layout and break circular dependencies
- Move error-handling, git-manager, prompt-manager, queue-validation, and reporting into src/services/ - Delete src/constants.ts — relocate AGENT_VALIDATORS and MCP_AGENT_MAPPING into session-manager.ts alongside agent definitions - Delete src/utils/output-formatter.ts — absorb filterJsonToolCalls and getAgentPrefix into ai/output-formatters.ts - Extract ActivityLogger interface into src/types/activity-logger.ts to break temporal/ → services circular dependency - Consolidate VulnType, ExploitationDecision into types/agents.ts and SessionMetadata into types/audit.ts - Remove dead timingResults/costResults globals from utils/metrics.ts and all consumers
This commit is contained in:
@@ -0,0 +1,325 @@
|
||||
// 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.
|
||||
|
||||
import { fs, path } from 'zx';
|
||||
import { PentestError } from './error-handling.js';
|
||||
import { ErrorCode } from '../types/errors.js';
|
||||
import { type Result, ok, err } from '../types/result.js';
|
||||
import { asyncPipe } from '../utils/functional.js';
|
||||
import type { VulnType, ExploitationDecision } from '../types/agents.js';
|
||||
|
||||
export type { VulnType, ExploitationDecision } from '../types/agents.js';
|
||||
|
||||
interface VulnTypeConfigItem {
|
||||
deliverable: string;
|
||||
queue: string;
|
||||
}
|
||||
|
||||
type VulnTypeConfig = Record<VulnType, VulnTypeConfigItem>;
|
||||
|
||||
type ErrorMessageResolver = string | ((existence: FileExistence) => string);
|
||||
|
||||
interface ValidationRule {
|
||||
predicate: (existence: FileExistence) => boolean;
|
||||
errorMessage: ErrorMessageResolver;
|
||||
retryable: boolean;
|
||||
}
|
||||
|
||||
interface FileExistence {
|
||||
deliverableExists: boolean;
|
||||
queueExists: boolean;
|
||||
}
|
||||
|
||||
interface PathsBase {
|
||||
vulnType: VulnType;
|
||||
deliverable: string;
|
||||
queue: string;
|
||||
sourceDir: string;
|
||||
}
|
||||
|
||||
interface PathsWithExistence extends PathsBase {
|
||||
existence: FileExistence;
|
||||
}
|
||||
|
||||
interface PathsWithQueue extends PathsWithExistence {
|
||||
queueData: QueueData;
|
||||
}
|
||||
|
||||
interface PathsWithError {
|
||||
error: PentestError;
|
||||
}
|
||||
|
||||
interface QueueData {
|
||||
vulnerabilities: unknown[];
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
interface QueueValidationResult {
|
||||
valid: boolean;
|
||||
data: QueueData | null;
|
||||
error: string | null;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Result type for safe validation - explicit error handling.
|
||||
*/
|
||||
export type SafeValidationResult = Result<ExploitationDecision, PentestError>;
|
||||
|
||||
// Vulnerability type configuration as immutable data
|
||||
const VULN_TYPE_CONFIG: VulnTypeConfig = Object.freeze({
|
||||
injection: Object.freeze({
|
||||
deliverable: 'injection_analysis_deliverable.md',
|
||||
queue: 'injection_exploitation_queue.json',
|
||||
}),
|
||||
xss: Object.freeze({
|
||||
deliverable: 'xss_analysis_deliverable.md',
|
||||
queue: 'xss_exploitation_queue.json',
|
||||
}),
|
||||
auth: Object.freeze({
|
||||
deliverable: 'auth_analysis_deliverable.md',
|
||||
queue: 'auth_exploitation_queue.json',
|
||||
}),
|
||||
ssrf: Object.freeze({
|
||||
deliverable: 'ssrf_analysis_deliverable.md',
|
||||
queue: 'ssrf_exploitation_queue.json',
|
||||
}),
|
||||
authz: Object.freeze({
|
||||
deliverable: 'authz_analysis_deliverable.md',
|
||||
queue: 'authz_exploitation_queue.json',
|
||||
}),
|
||||
}) as VulnTypeConfig;
|
||||
|
||||
// Pure function to create validation rule
|
||||
function createValidationRule(
|
||||
predicate: (existence: FileExistence) => boolean,
|
||||
errorMessage: ErrorMessageResolver,
|
||||
retryable: boolean = true
|
||||
): ValidationRule {
|
||||
return Object.freeze({ predicate, errorMessage, retryable });
|
||||
}
|
||||
|
||||
// Symmetric deliverable rules: queue and deliverable must exist together (prevents partial analysis from triggering exploitation)
|
||||
const fileExistenceRules: readonly ValidationRule[] = Object.freeze([
|
||||
createValidationRule(
|
||||
({ deliverableExists, queueExists }) => deliverableExists && queueExists,
|
||||
getExistenceErrorMessage
|
||||
),
|
||||
]);
|
||||
|
||||
// Generate appropriate error message based on which files are missing
|
||||
function getExistenceErrorMessage(existence: FileExistence): string {
|
||||
const { deliverableExists, queueExists } = existence;
|
||||
|
||||
if (!deliverableExists && !queueExists) {
|
||||
return 'Analysis failed: Neither deliverable nor queue file exists. Analysis agent must create both files.';
|
||||
}
|
||||
if (!queueExists) {
|
||||
return 'Analysis incomplete: Deliverable exists but queue file missing. Analysis agent must create both files.';
|
||||
}
|
||||
return 'Analysis incomplete: Queue exists but deliverable file missing. Analysis agent must create both files.';
|
||||
}
|
||||
|
||||
// Pure function to create file paths
|
||||
const createPaths = (
|
||||
vulnType: VulnType,
|
||||
sourceDir: string
|
||||
): PathsBase | PathsWithError => {
|
||||
const config = VULN_TYPE_CONFIG[vulnType];
|
||||
if (!config) {
|
||||
return {
|
||||
error: new PentestError(
|
||||
`Unknown vulnerability type: ${vulnType}`,
|
||||
'validation',
|
||||
false,
|
||||
{ vulnType }
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
return Object.freeze({
|
||||
vulnType,
|
||||
deliverable: path.join(sourceDir, 'deliverables', config.deliverable),
|
||||
queue: path.join(sourceDir, 'deliverables', config.queue),
|
||||
sourceDir,
|
||||
});
|
||||
};
|
||||
|
||||
// Pure function to check file existence
|
||||
const checkFileExistence = async (
|
||||
paths: PathsBase | PathsWithError
|
||||
): Promise<PathsWithExistence | PathsWithError> => {
|
||||
if ('error' in paths) return paths;
|
||||
|
||||
const [deliverableExists, queueExists] = await Promise.all([
|
||||
fs.pathExists(paths.deliverable),
|
||||
fs.pathExists(paths.queue),
|
||||
]);
|
||||
|
||||
return Object.freeze({
|
||||
...paths,
|
||||
existence: Object.freeze({ deliverableExists, queueExists }),
|
||||
});
|
||||
};
|
||||
|
||||
// Validates deliverable/queue symmetry - both must exist or neither
|
||||
const validateExistenceRules = (
|
||||
pathsWithExistence: PathsWithExistence | PathsWithError
|
||||
): PathsWithExistence | PathsWithError => {
|
||||
if ('error' in pathsWithExistence) return pathsWithExistence;
|
||||
|
||||
const { existence, vulnType } = pathsWithExistence;
|
||||
|
||||
// Find the first rule that fails
|
||||
const failedRule = fileExistenceRules.find((rule) => !rule.predicate(existence));
|
||||
|
||||
if (failedRule) {
|
||||
const message =
|
||||
typeof failedRule.errorMessage === 'function'
|
||||
? failedRule.errorMessage(existence)
|
||||
: failedRule.errorMessage;
|
||||
|
||||
return {
|
||||
error: new PentestError(
|
||||
`${message} (${vulnType})`,
|
||||
'validation',
|
||||
failedRule.retryable,
|
||||
{
|
||||
vulnType,
|
||||
deliverablePath: pathsWithExistence.deliverable,
|
||||
queuePath: pathsWithExistence.queue,
|
||||
existence,
|
||||
},
|
||||
ErrorCode.DELIVERABLE_NOT_FOUND
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
return pathsWithExistence;
|
||||
};
|
||||
|
||||
// Pure function to validate queue structure
|
||||
const validateQueueStructure = (content: string): QueueValidationResult => {
|
||||
try {
|
||||
const parsed = JSON.parse(content) as unknown;
|
||||
const isValid =
|
||||
typeof parsed === 'object' &&
|
||||
parsed !== null &&
|
||||
'vulnerabilities' in parsed &&
|
||||
Array.isArray((parsed as QueueData).vulnerabilities);
|
||||
|
||||
return Object.freeze({
|
||||
valid: isValid,
|
||||
data: isValid ? (parsed as QueueData) : null,
|
||||
error: null,
|
||||
});
|
||||
} catch (parseError) {
|
||||
return Object.freeze({
|
||||
valid: false,
|
||||
data: null,
|
||||
error: parseError instanceof Error ? parseError.message : String(parseError),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
// Queue parse failures are retryable - agent can fix malformed JSON on retry
|
||||
const validateQueueContent = async (
|
||||
pathsWithExistence: PathsWithExistence | PathsWithError
|
||||
): Promise<PathsWithQueue | PathsWithError> => {
|
||||
if ('error' in pathsWithExistence) return pathsWithExistence;
|
||||
|
||||
try {
|
||||
const queueContent = await fs.readFile(pathsWithExistence.queue, 'utf8');
|
||||
const queueValidation = validateQueueStructure(queueContent);
|
||||
|
||||
if (!queueValidation.valid) {
|
||||
// Rule 6: Both exist, queue invalid
|
||||
return {
|
||||
error: new PentestError(
|
||||
queueValidation.error
|
||||
? `Queue validation failed for ${pathsWithExistence.vulnType}: Invalid JSON structure. Analysis agent must fix queue format.`
|
||||
: `Queue validation failed for ${pathsWithExistence.vulnType}: Missing or invalid 'vulnerabilities' array. Analysis agent must fix queue structure.`,
|
||||
'validation',
|
||||
true, // retryable
|
||||
{
|
||||
vulnType: pathsWithExistence.vulnType,
|
||||
queuePath: pathsWithExistence.queue,
|
||||
originalError: queueValidation.error,
|
||||
queueStructure: queueValidation.data ? Object.keys(queueValidation.data) : [],
|
||||
}
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
return Object.freeze({
|
||||
...pathsWithExistence,
|
||||
queueData: queueValidation.data!,
|
||||
});
|
||||
} catch (readError) {
|
||||
return {
|
||||
error: new PentestError(
|
||||
`Failed to read queue file for ${pathsWithExistence.vulnType}: ${readError instanceof Error ? readError.message : String(readError)}`,
|
||||
'filesystem',
|
||||
false,
|
||||
{
|
||||
vulnType: pathsWithExistence.vulnType,
|
||||
queuePath: pathsWithExistence.queue,
|
||||
originalError: readError instanceof Error ? readError.message : String(readError),
|
||||
}
|
||||
),
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
// Final decision: skip if queue says no vulns, proceed if vulns found, error otherwise
|
||||
const determineExploitationDecision = (
|
||||
validatedData: PathsWithQueue | PathsWithError
|
||||
): ExploitationDecision => {
|
||||
if ('error' in validatedData) {
|
||||
throw validatedData.error;
|
||||
}
|
||||
|
||||
const hasVulnerabilities = validatedData.queueData.vulnerabilities.length > 0;
|
||||
|
||||
// Rule 4: Both exist, queue valid and populated
|
||||
// Rule 5: Both exist, queue valid but empty
|
||||
return Object.freeze({
|
||||
shouldExploit: hasVulnerabilities,
|
||||
shouldRetry: false,
|
||||
vulnerabilityCount: validatedData.queueData.vulnerabilities.length,
|
||||
vulnType: validatedData.vulnType,
|
||||
});
|
||||
};
|
||||
|
||||
// Main functional validation pipeline
|
||||
export async function validateQueueAndDeliverable(
|
||||
vulnType: VulnType,
|
||||
sourceDir: string
|
||||
): Promise<ExploitationDecision> {
|
||||
return asyncPipe<ExploitationDecision>(
|
||||
createPaths(vulnType, sourceDir),
|
||||
checkFileExistence,
|
||||
validateExistenceRules,
|
||||
validateQueueContent,
|
||||
determineExploitationDecision
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Safely validate queue and deliverable files.
|
||||
* Returns Result<ExploitationDecision, PentestError> for explicit error handling.
|
||||
*/
|
||||
export async function validateQueueSafe(
|
||||
vulnType: VulnType,
|
||||
sourceDir: string
|
||||
): Promise<SafeValidationResult> {
|
||||
try {
|
||||
const result = await validateQueueAndDeliverable(vulnType, sourceDir);
|
||||
return ok(result);
|
||||
} catch (error) {
|
||||
return err(error as PentestError);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user