serverMCP.js•9.59 kB
#!/usr/bin/env node
/**
* MalwareAnalyzerMCP
*
* MCP server for Claude Desktop that allows executing terminal commands for malware analysis.
* Implements the Model Context Protocol with tools for command execution and output reading.
*/
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { zodToJsonSchema } from 'zod-to-json-schema';
import { z } from 'zod';
import { TerminalManager } from './terminalManager.js';
import { commands } from './commands.js';
/**
* Custom stdio transport with filtering capability
* Handles communication between Claude Desktop and our server
*/
class FilteredStdioTransport extends StdioServerTransport {
constructor() {
// Create a proxy for stdout that only allows valid JSON
const originalStdoutWrite = process.stdout.write;
process.stdout.write = function(buffer) {
// Only intercept string output that doesn't look like JSON
if (typeof buffer === 'string' && !buffer.trim().startsWith('{')) {
return process.stderr.write(buffer);
}
return originalStdoutWrite.apply(process.stdout, arguments);
};
super();
// Log initialization to stderr
process.stderr.write(`[malware-analyzer] Initialized FilteredStdioTransport\n`);
}
}
/**
* Schema for shellCommand tool
* Defines parameters for executing terminal commands
*/
const shellCommandSchema = z.object({
command: z.string().min(1).describe("The command to execute in the terminal"),
timeout_ms: z.number().optional().describe("Optional timeout in milliseconds (default: 30000)")
});
/**
* Schema for read_output tool
* Defines parameters for reading process output
*/
const readOutputSchema = z.object({
pid: z.number().int().describe("The process ID to read output from")
});
/**
* Main entry point for the MCP server
* Configures and starts the server with tool handlers
*/
async function main() {
try {
// Setup error handling
process.on('uncaughtException', (error) => {
console.error(`Uncaught exception: ${error instanceof Error ? error.message : String(error)}`);
process.exit(1);
});
process.on('unhandledRejection', (reason) => {
console.error(`Unhandled rejection: ${reason instanceof Error ? reason.message : String(reason)}`);
process.exit(1);
});
console.error("Initializing MalwareAnalyzerMCP server...");
// Initialize transport
const transport = new FilteredStdioTransport();
// Initialize terminal manager for process handling
const terminalManager = new TerminalManager();
// Create server instance
const server = new Server(
{
name: "malware-analyzer",
version: "1.0.0",
},
{
capabilities: {
tools: {}, // Enable tools capability
},
}
);
// Configure tools list handler
server.setRequestHandler(ListToolsRequestSchema, async () => {
// List of basic tools
const basicTools = [
{
name: 'shell_command',
description: 'Execute a command in the terminal with timeout. Command will continue running in background if it doesn\'t complete within timeout.',
inputSchema: zodToJsonSchema(shellCommandSchema),
},
{
name: 'read_output',
description: 'Read output from a running or completed process.',
inputSchema: zodToJsonSchema(readOutputSchema),
},
];
// Generate tools from commands configuration
const specializedTools = Object.values(commands).map(cmd => ({
name: cmd.name,
description: cmd.description + (cmd.helpText ? '\n' + cmd.helpText : ''),
inputSchema: zodToJsonSchema(cmd.schema),
}));
return {
tools: [...basicTools, ...specializedTools],
};
});
// Configure tool execution handler
server.setRequestHandler(CallToolRequestSchema, async (request) => {
try {
const { name, arguments: args } = request.params;
// Check if this is a specialized command
if (commands[name]) {
try {
const cmdConfig = commands[name];
// Validate arguments against schema
const validationResult = cmdConfig.schema.safeParse(args);
if (!validationResult.success) {
return {
content: [{
type: "text",
text: `Error: Invalid parameters for ${name} command.\n${JSON.stringify(validationResult.error.format())}`
}],
isError: true,
};
}
// Build the command string
const commandStr = cmdConfig.buildCommand(validationResult.data);
console.error(`Executing specialized command: ${commandStr}`);
// Execute the command via the terminal manager
const result = await terminalManager.shellCommand(commandStr);
console.error(`${name} command executed with PID: ${result.pid}, blocked: ${result.isBlocked}`);
return {
content: [{ type: "text", text: JSON.stringify(result) }],
};
} catch (error) {
console.error(`Error executing ${name} command:`, error);
return {
content: [{ type: "text", text: `Error: ${error instanceof Error ? error.message : String(error)}` }],
isError: true,
};
}
}
// Handle basic tools
switch (name) {
case 'shell_command':
try {
// Type-check and validate arguments
if (!args || typeof args.command !== 'string') {
return {
content: [{ type: "text", text: "Error: Invalid command parameter" }],
isError: true,
};
}
console.error(`Executing command: ${args.command}`);
const result = await terminalManager.shellCommand(
args.command,
typeof args.timeout_ms === 'number' ? args.timeout_ms : undefined
);
console.error(`Command executed with PID: ${result.pid}, blocked: ${result.isBlocked}`);
return {
content: [{ type: "text", text: JSON.stringify(result) }],
};
} catch (error) {
console.error('Error executing command:', error);
return {
content: [{ type: "text", text: `Error: ${error instanceof Error ? error.message : String(error)}` }],
isError: true,
};
}
case 'read_output':
try {
// Type-check and validate arguments
if (!args || typeof args.pid !== 'number') {
return {
content: [{ type: "text", text: "Error: Invalid PID parameter" }],
isError: true,
};
}
console.error(`Reading output for PID: ${args.pid}`);
const result = terminalManager.readOutput(args.pid);
return {
content: [{ type: "text", text: JSON.stringify(result) }],
};
} catch (error) {
console.error('Error reading output:', error);
return {
content: [{ type: "text", text: `Error: ${error instanceof Error ? error.message : String(error)}` }],
isError: true,
};
}
default:
return {
content: [{ type: "text", text: `Error: Unknown tool: ${name}` }],
isError: true,
};
}
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
return {
content: [{ type: "text", text: `Error: ${errorMessage}` }],
isError: true,
};
}
});
// Setup graceful shutdown handler
const cleanup = () => {
console.error('Initiating server cleanup...');
terminalManager.shutdown();
server.close().then(() => {
console.error('Server closed.');
process.exit(0);
}).catch((err) => {
console.error('Error closing server:', err);
process.exit(1);
});
// Force exit after a timeout if cleanup hangs
setTimeout(() => {
console.error('Cleanup timed out. Forcing exit.');
process.exit(1);
}, 5000);
};
// Register shutdown handlers
process.on('SIGINT', cleanup);
process.on('SIGTERM', cleanup);
// Connect server to transport
console.error('Starting MalwareAnalyzerMCP server...');
await server.connect(transport);
console.error('Server started and listening.');
} catch (error) {
console.error('Fatal error during server startup:');
console.error(error instanceof Error ? error.message : String(error));
console.error(error instanceof Error && error.stack ? error.stack : 'No stack trace available');
process.exit(1);
}
}
// Start the server
main().catch(error => {
console.error('Fatal error during server initialization:');
console.error(error instanceof Error ? error.message : String(error));
process.exit(1);
});