tot-engine.js•2.64 kB
export class ToTEngine {
constructor() {
this.beamWidth = 5;
this.maxDepth = 3;
}
// Pentest-specific scoring: prioritize high-value assets and high-severity vulns
async evaluateAttackStep(attackStep) {
// Example: prioritize steps mentioning 'domain controller', 'database', or 'CVE'
let score = 0.5;
if (/domain controller|database/i.test(attackStep.description)) score += 0.3;
if (/CVE-\d{4}-\d{4,7}/i.test(attackStep.description)) score += 0.2;
if (/critical|high/i.test(attackStep.description)) score += 0.2;
// Add randomness for diversity
score += Math.random() * 0.2;
return Math.min(score, 1.0);
}
// Generate attack steps with asset tagging
async generateAttackSteps(context, numSteps = 3) {
const attackSteps = [];
const scores = [];
const assets = ['10.0.0.5', 'webserver', 'database']; // Example assets
for (let i = 0; i < numSteps; i++) {
// Simulate attack step generation
const asset = assets[i % assets.length];
const description = `Attack Step ${i + 1} on ${asset} about ${context}`;
const attackStep = { description, asset };
const score = await this.evaluateAttackStep(attackStep);
attackSteps.push(attackStep);
scores.push(score);
}
return { attackSteps, scores };
}
async expandPath(path) {
const lastStep = path[path.length - 1];
const { attackSteps, scores } = await this.generateAttackSteps(lastStep.description || lastStep);
return attackSteps.map((attackStep, i) => ({
path: [...path, attackStep],
score: scores[i],
asset: attackStep.asset,
recommendedTool: this.recommendTool(attackStep)
}));
}
// Recommend a tool based on the attack step
recommendTool(attackStep) {
if (/scan|nmap/i.test(attackStep.description)) return 'nmap';
if (/exploit|CVE/i.test(attackStep.description)) return 'metasploit';
if (/database/i.test(attackStep.description)) return 'sqlmap';
return 'manual investigation';
}
async search(initialContext) {
let paths = [{ path: [{ description: initialContext, asset: null }], score: 1.0 }];
for (let depth = 0; depth < this.maxDepth; depth++) {
const newPaths = [];
for (const currentPath of paths) {
const expansions = await this.expandPath(currentPath.path);
newPaths.push(...expansions);
}
// Keep top k paths based on scores
paths = newPaths
.sort((a, b) => b.score - a.score)
.slice(0, this.beamWidth);
}
// Highlight the most critical path
if (paths.length > 0) paths[0].critical = true;
return paths;
}
}