stellar-cyber-mcp-agents
Version:
Model Context Protocol (MCP) server for Stellar Cyber security operations with specialized multi-agent analysis capabilities
561 lines • 21.9 kB
JavaScript
import { BaseAgent } from '../core/base-agent.js';
import { AgentType, AgentHealth } from '../types/agent.js';
export class WorkflowAgent extends BaseAgent {
agentId;
accessToken = null;
tokenExpiresAt = 0;
config;
workflows = new Map();
executions = new Map();
agentRegistry = new Map();
constructor(metadata, registry, channel, logger, metrics, config) {
super(metadata, registry, channel, logger, metrics);
this.agentId = {
type: AgentType.WORKFLOW,
instance: 'primary',
uuid: crypto.randomUUID()
};
this.config = {
apiUrl: config.apiUrl,
apiToken: config.apiToken,
maxConcurrentWorkflows: config.maxConcurrentWorkflows || 5,
defaultTimeout: config.defaultTimeout || 300000
};
this.initializeDefaultWorkflows();
}
getAgentId() {
return this.agentId;
}
// Removed duplicate onStart and onStop - implemented at bottom of class
async performHealthCheck() {
try {
if (!this.accessToken || Date.now() >= this.tokenExpiresAt - 30000) {
await this.performTokenRefresh();
}
const response = await fetch(`${this.config.apiUrl}/connect/api/v1/health`, {
method: 'GET',
headers: {
'Authorization': `Bearer ${this.accessToken}`,
'Content-Type': 'application/json'
}
});
return response.ok;
}
catch (error) {
console.error('Health check failed:', error);
return false;
}
}
async performTokenRefresh() {
try {
const response = await fetch(`${this.config.apiUrl}/connect/api/v1/access_token`, {
method: 'POST',
headers: {
'Authorization': `Bearer ${this.config.apiToken}`,
'Content-Type': 'application/json'
}
});
if (!response.ok) {
throw new Error(`Token refresh failed: ${response.status}`);
}
const data = await response.json();
this.accessToken = data.access_token;
this.tokenExpiresAt = Date.now() + (data.exp * 1000);
}
catch (error) {
console.error('Token refresh failed:', error);
throw error;
}
}
initializeDefaultWorkflows() {
// Comprehensive Threat Investigation Workflow
const threatInvestigationWorkflow = {
id: 'threat-investigation-comprehensive',
name: 'Comprehensive Threat Investigation',
description: 'Full investigation workflow including all analysis types',
version: '1.0.0',
trigger: {
type: 'manual'
},
steps: [
{
id: 'step-1',
name: 'Initial Investigation',
agentType: AgentType.INVESTIGATION,
action: 'investigate_case',
parameters: {},
dependencies: [],
timeout: 60000,
retryCount: 2,
optional: false
},
{
id: 'step-2',
name: 'Network Analysis',
agentType: AgentType.NETWORK_ANALYSIS,
action: 'analyze_network',
parameters: {},
dependencies: ['step-1'],
timeout: 120000,
retryCount: 1,
optional: false
},
{
id: 'step-3',
name: 'Malware Analysis',
agentType: AgentType.MALWARE_ANALYSIS,
action: 'analyze_malware',
parameters: {},
dependencies: ['step-1'],
timeout: 180000,
retryCount: 1,
optional: true
},
{
id: 'step-4',
name: 'Credential Analysis',
agentType: AgentType.CREDENTIAL_ANALYSIS,
action: 'analyze_credentials',
parameters: {},
dependencies: ['step-1'],
timeout: 90000,
retryCount: 1,
optional: true
},
{
id: 'step-5',
name: 'Correlation Analysis',
agentType: AgentType.CORRELATION,
action: 'find_related_cases',
parameters: {},
dependencies: ['step-1', 'step-2'],
timeout: 120000,
retryCount: 1,
optional: false
}
],
metadata: {
created: new Date().toISOString(),
author: 'system',
tags: ['investigation', 'comprehensive', 'threat-analysis'],
category: 'Security Investigation'
}
};
// Quick Triage Workflow
const quickTriageWorkflow = {
id: 'quick-triage',
name: 'Quick Threat Triage',
description: 'Fast initial triage for incoming cases',
version: '1.0.0',
trigger: {
type: 'event'
},
steps: [
{
id: 'step-1',
name: 'Basic Investigation',
agentType: AgentType.INVESTIGATION,
action: 'investigate_case',
parameters: { quick_mode: true },
dependencies: [],
timeout: 30000,
retryCount: 1,
optional: false
},
{
id: 'step-2',
name: 'Risk Assessment',
agentType: AgentType.CORRELATION,
action: 'assess_risk',
parameters: { basic_mode: true },
dependencies: ['step-1'],
timeout: 30000,
retryCount: 1,
optional: false,
condition: {
field: 'step-1.summary.riskScore',
operator: 'greater_than',
value: 50
}
}
],
metadata: {
created: new Date().toISOString(),
author: 'system',
tags: ['triage', 'quick', 'initial-assessment'],
category: 'Security Triage'
}
};
// Malware Analysis Workflow
const malwareAnalysisWorkflow = {
id: 'malware-deep-analysis',
name: 'Deep Malware Analysis',
description: 'Comprehensive malware analysis and artifact examination',
version: '1.0.0',
trigger: {
type: 'manual'
},
steps: [
{
id: 'step-1',
name: 'Basic Investigation',
agentType: AgentType.INVESTIGATION,
action: 'investigate_case',
parameters: {},
dependencies: [],
timeout: 60000,
retryCount: 2,
optional: false
},
{
id: 'step-2',
name: 'Malware Analysis',
agentType: AgentType.MALWARE_ANALYSIS,
action: 'analyze_malware',
parameters: { detailed: true },
dependencies: ['step-1'],
timeout: 300000,
retryCount: 1,
optional: false
},
{
id: 'step-3',
name: 'Network IOC Analysis',
agentType: AgentType.NETWORK_ANALYSIS,
action: 'analyze_network',
parameters: { focus: 'iocs' },
dependencies: ['step-2'],
timeout: 120000,
retryCount: 1,
optional: false
},
{
id: 'step-4',
name: 'Related Case Search',
agentType: AgentType.CORRELATION,
action: 'find_related_cases',
parameters: { malware_focus: true },
dependencies: ['step-2'],
timeout: 90000,
retryCount: 1,
optional: true
}
],
metadata: {
created: new Date().toISOString(),
author: 'system',
tags: ['malware', 'deep-analysis', 'forensics'],
category: 'Malware Analysis'
}
};
this.workflows.set(threatInvestigationWorkflow.id, threatInvestigationWorkflow);
this.workflows.set(quickTriageWorkflow.id, quickTriageWorkflow);
this.workflows.set(malwareAnalysisWorkflow.id, malwareAnalysisWorkflow);
}
registerAgent(agentType, agentInstance) {
this.agentRegistry.set(agentType, agentInstance);
}
async executeWorkflow(request) {
const workflow = this.workflows.get(request.workflowId);
if (!workflow) {
throw new Error(`Workflow not found: ${request.workflowId}`);
}
const executionId = crypto.randomUUID();
const execution = {
id: executionId,
workflowId: request.workflowId,
caseId: request.caseId,
status: 'PENDING',
startTime: new Date().toISOString(),
results: new Map(),
errors: [],
metrics: {
totalSteps: workflow.steps.length,
completedSteps: 0,
failedSteps: 0
}
};
this.executions.set(executionId, execution);
// Start execution asynchronously
this.executeWorkflowSteps(execution, workflow, request.parameters || {})
.catch(error => {
execution.status = 'FAILED';
execution.endTime = new Date().toISOString();
execution.errors.push({
stepId: 'workflow',
error: error.message,
timestamp: new Date().toISOString()
});
});
this.emit('workflow:started', {
executionId,
workflowId: request.workflowId,
caseId: request.caseId
});
return executionId;
}
async executeWorkflowSteps(execution, workflow, parameters) {
execution.status = 'RUNNING';
const completedSteps = new Set();
try {
// Execute steps in dependency order
while (completedSteps.size < workflow.steps.length) {
const readySteps = workflow.steps.filter(step => !completedSteps.has(step.id) &&
step.dependencies.every(dep => completedSteps.has(dep)));
if (readySteps.length === 0) {
// Check if there are remaining steps that can't be executed
const remainingSteps = workflow.steps.filter(step => !completedSteps.has(step.id));
if (remainingSteps.length > 0) {
throw new Error('Workflow has circular dependencies or missing steps');
}
break;
}
// Execute ready steps in parallel
const stepPromises = readySteps.map(step => this.executeStep(execution, step, parameters));
const stepResults = await Promise.allSettled(stepPromises);
for (let i = 0; i < readySteps.length; i++) {
const step = readySteps[i];
const result = stepResults[i];
if (result.status === 'fulfilled') {
execution.results.set(step.id, result.value);
execution.metrics.completedSteps++;
completedSteps.add(step.id);
}
else if (!step.optional) {
execution.metrics.failedSteps++;
execution.errors.push({
stepId: step.id,
error: result.reason?.message || 'Step execution failed',
timestamp: new Date().toISOString()
});
throw new Error(`Required step failed: ${step.id}`);
}
else {
// Optional step failed, continue
execution.metrics.failedSteps++;
execution.errors.push({
stepId: step.id,
error: result.reason?.message || 'Optional step failed',
timestamp: new Date().toISOString()
});
completedSteps.add(step.id);
}
}
}
execution.status = 'COMPLETED';
execution.endTime = new Date().toISOString();
execution.metrics.duration = new Date(execution.endTime).getTime() -
new Date(execution.startTime).getTime();
this.emit('workflow:completed', {
executionId: execution.id,
workflowId: execution.workflowId,
duration: execution.metrics.duration,
results: Object.fromEntries(execution.results)
});
}
catch (error) {
execution.status = 'FAILED';
execution.endTime = new Date().toISOString();
execution.metrics.duration = new Date(execution.endTime).getTime() -
new Date(execution.startTime).getTime();
this.emit('workflow:failed', {
executionId: execution.id,
workflowId: execution.workflowId,
error: error instanceof Error ? error.message : String(error),
duration: execution.metrics.duration
});
throw error;
}
}
async executeStep(execution, step, globalParameters) {
execution.currentStep = step.id;
// Check step condition if present
if (step.condition && !this.evaluateCondition(step.condition, execution.results)) {
return { skipped: true, reason: 'Condition not met' };
}
// Get the agent for this step
const agent = this.agentRegistry.get(step.agentType);
if (!agent) {
throw new Error(`Agent not found for type: ${step.agentType}`);
}
// Prepare step parameters
const stepParameters = {
...globalParameters,
...step.parameters,
caseId: execution.caseId
};
// Execute the step with retry logic
let lastError = null;
for (let attempt = 0; attempt <= step.retryCount; attempt++) {
try {
const timeout = new Promise((_, reject) => {
setTimeout(() => reject(new Error('Step timeout')), step.timeout);
});
const stepPromise = this.invokeAgentAction(agent, step.action, stepParameters);
const result = await Promise.race([stepPromise, timeout]);
this.emit('workflow:step_completed', {
executionId: execution.id,
stepId: step.id,
attempt: attempt + 1,
result
});
return result;
}
catch (error) {
lastError = error;
if (attempt < step.retryCount) {
await this.delay(1000 * (attempt + 1)); // Exponential backoff
}
}
}
throw lastError || new Error('Step failed after retries');
}
evaluateCondition(condition, results) {
const fieldPath = condition.field.split('.');
let value = Object.fromEntries(results);
// Navigate to the field value
for (const path of fieldPath) {
if (value && typeof value === 'object' && path in value) {
value = value[path];
}
else {
return false; // Field not found
}
}
// Evaluate condition
switch (condition.operator) {
case 'equals':
return value === condition.value;
case 'greater_than':
return typeof value === 'number' && value > condition.value;
case 'less_than':
return typeof value === 'number' && value < condition.value;
case 'contains':
return typeof value === 'string' && value.includes(condition.value);
case 'exists':
return value !== undefined && value !== null;
default:
return false;
}
}
async invokeAgentAction(agent, action, parameters) {
// Map action names to agent methods
const actionMap = {
'investigate_case': 'investigateCase',
'analyze_network': 'analyzeNetwork',
'analyze_malware': 'analyzeMalware',
'analyze_credentials': 'analyzeCredentials',
'find_related_cases': 'findRelatedCases',
'assess_risk': 'assessRisk'
};
const methodName = actionMap[action];
if (!methodName || typeof agent[methodName] !== 'function') {
throw new Error(`Action not supported: ${action}`);
}
return await agent[methodName](parameters);
}
delay(ms) {
return new Promise(resolve => setTimeout(resolve, ms));
}
async getWorkflowStatus(executionId) {
return this.executions.get(executionId) || null;
}
async cancelWorkflow(executionId) {
const execution = this.executions.get(executionId);
if (!execution || execution.status !== 'RUNNING') {
return false;
}
execution.status = 'CANCELLED';
execution.endTime = new Date().toISOString();
execution.metrics.duration = new Date(execution.endTime).getTime() -
new Date(execution.startTime).getTime();
this.emit('workflow:cancelled', {
executionId,
workflowId: execution.workflowId
});
return true;
}
getAvailableWorkflows() {
return Array.from(this.workflows.values());
}
async addWorkflow(workflow) {
// Validate workflow
this.validateWorkflow(workflow);
this.workflows.set(workflow.id, workflow);
this.emit('workflow:added', {
workflowId: workflow.id,
name: workflow.name
});
}
validateWorkflow(workflow) {
if (!workflow.id || !workflow.name || !workflow.steps) {
throw new Error('Invalid workflow: missing required fields');
}
// Check for circular dependencies
const visited = new Set();
const recursionStack = new Set();
const hasCycle = (stepId) => {
if (recursionStack.has(stepId)) {
return true;
}
if (visited.has(stepId)) {
return false;
}
visited.add(stepId);
recursionStack.add(stepId);
const step = workflow.steps.find(s => s.id === stepId);
if (step) {
for (const dep of step.dependencies) {
if (hasCycle(dep)) {
return true;
}
}
}
recursionStack.delete(stepId);
return false;
};
for (const step of workflow.steps) {
if (hasCycle(step.id)) {
throw new Error('Workflow contains circular dependencies');
}
}
}
async getExecutionHistory(workflowId, limit = 50) {
let executions = Array.from(this.executions.values());
if (workflowId) {
executions = executions.filter(e => e.workflowId === workflowId);
}
return executions
.sort((a, b) => new Date(b.startTime).getTime() - new Date(a.startTime).getTime())
.slice(0, limit);
}
// Required abstract method implementations from BaseAgent
async onInitialize() {
this.logger.info('Workflow Agent initialized');
}
async onStart() {
this.logger.info('Workflow Agent started');
}
async onStop() {
this.logger.info('Workflow Agent stopped');
}
async onDestroy() {
this.logger.info('Workflow Agent destroyed');
}
async onHealthCheck() {
return AgentHealth.HEALTHY;
}
async handleRequest(request, context) {
switch (request.capability) {
case 'execute_workflow':
return await this.executeWorkflow(request.payload);
case 'get_workflow_status':
return await this.getWorkflowStatus(request.payload.executionId);
case 'cancel_workflow':
return await this.cancelWorkflow(request.payload.executionId);
default:
throw new Error(`Unsupported capability: ${request.capability}`);
}
}
}
//# sourceMappingURL=workflow-agent.js.map