index.ts

#!/usr/bin/env node

/**
 * ExploitDB MCP Server
 *
 * Developed by Cyreslab.ai (https://cyreslab.ai)
 * Contact: contact@cyreslab.ai
 * GitHub: https://github.com/Cyreslab-AI
 *
 * This server provides access to ExploitDB functionality through the Model Context Protocol.
 * It allows AI assistants to query information about security exploits and vulnerabilities,
 * enhancing cybersecurity research and threat intelligence capabilities.
 *
 * Copyright (c) 2025 Cyreslab.ai. All rights reserved.
 * Licensed under the MIT License.
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  ListResourcesRequestSchema,
  ListToolsRequestSchema,
  McpError,
  ReadResourceRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';
import config, { ensureDataDir } from './config.js';
import db from './db/index.js';
import exploitdb from './exploitdb/index.js';

/**
 * ExploitDB MCP Server class
 */
class ExploitDBServer {
  private server: Server;
  private updateInterval: NodeJS.Timeout | null = null;

  constructor() {
    this.server = new Server(
      {
        name: 'mcp-exploitdb-server',
        version: '0.1.0',
      },
      {
        capabilities: {
          resources: {},
          tools: {},
        },
      }
    );

    this.setupResourceHandlers();
    this.setupToolHandlers();

    // Error handling
    this.server.onerror = (error) => console.error('[MCP Error]', error);
    process.on('SIGINT', async () => {
      await this.close();
      process.exit(0);
    });
  }

  /**
   * Set up resource handlers
   * For now, we're not implementing any static resources
   */
  private setupResourceHandlers() {
    this.server.setRequestHandler(ListResourcesRequestSchema, async () => {
      return {
        resources: []
      };
    });

    this.server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
      throw new McpError(
        ErrorCode.InvalidRequest,
        `Invalid URI: ${request.params.uri}`
      );
    });
  }

  /**
   * Set up tool handlers for ExploitDB functionality
   */
  private setupToolHandlers() {
    this.server.setRequestHandler(ListToolsRequestSchema, async () => {
      return {
        tools: [
          {
            name: 'search_exploits',
            description: 'Search for exploits in the ExploitDB database',
            inputSchema: {
              type: 'object',
              properties: {
                query: {
                  type: 'string',
                  description: 'Search query (searches in description and CVE)'
                },
                platform: {
                  type: 'string',
                  description: 'Filter by platform (e.g., windows, linux, php, etc.)'
                },
                type: {
                  type: 'string',
                  description: 'Filter by exploit type (e.g., webapps, remote, local, etc.)'
                },
                cve: {
                  type: 'string',
                  description: 'Filter by CVE ID (e.g., CVE-2021-44228)'
                },
                author: {
                  type: 'string',
                  description: 'Filter by author name'
                },
                start_date: {
                  type: 'string',
                  description: 'Filter by start date (YYYY-MM-DD)'
                },
                end_date: {
                  type: 'string',
                  description: 'Filter by end date (YYYY-MM-DD)'
                },
                verified: {
                  type: 'boolean',
                  description: 'Filter by verified status'
                },
                limit: {
                  type: 'number',
                  description: `Maximum number of results to return (default: ${config.maxResults})`
                },
                offset: {
                  type: 'number',
                  description: 'Offset for pagination (default: 0)'
                }
              }
            }
          },
          {
            name: 'get_exploit',
            description: 'Get detailed information about a specific exploit',
            inputSchema: {
              type: 'object',
              properties: {
                id: {
                  type: 'number',
                  description: 'Exploit ID'
                },
                include_code: {
                  type: 'boolean',
                  description: 'Whether to include the exploit code (default: true)'
                }
              },
              required: ['id']
            }
          },
          {
            name: 'find_by_cve',
            description: 'Find exploits by CVE ID',
            inputSchema: {
              type: 'object',
              properties: {
                cve: {
                  type: 'string',
                  description: 'CVE ID (e.g., CVE-2021-44228)'
                },
                limit: {
                  type: 'number',
                  description: `Maximum number of results to return (default: ${config.maxResults})`
                }
              },
              required: ['cve']
            }
          },
          {
            name: 'get_recent_exploits',
            description: 'Get recently added exploits',
            inputSchema: {
              type: 'object',
              properties: {
                limit: {
                  type: 'number',
                  description: `Maximum number of results to return (default: ${config.maxResults})`
                }
              }
            }
          },
          {
            name: 'get_statistics',
            description: 'Get statistics about the exploits in the database',
            inputSchema: {
              type: 'object',
              properties: {}
            }
          }
        ]
      };
    });

    this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
      switch (request.params.name) {
        case 'search_exploits': {
          const query = String(request.params.arguments?.query || '');
          const platform = request.params.arguments?.platform ? String(request.params.arguments.platform) : undefined;
          const type = request.params.arguments?.type ? String(request.params.arguments.type) : undefined;
          const cve = request.params.arguments?.cve ? String(request.params.arguments.cve) : undefined;
          const author = request.params.arguments?.author ? String(request.params.arguments.author) : undefined;
          const startDate = request.params.arguments?.start_date ? String(request.params.arguments.start_date) : undefined;
          const endDate = request.params.arguments?.end_date ? String(request.params.arguments.end_date) : undefined;
          const verified = request.params.arguments?.verified !== undefined ? Boolean(request.params.arguments.verified) : undefined;
          const limit = request.params.arguments?.limit ? Number(request.params.arguments.limit) : config.maxResults;
          const offset = request.params.arguments?.offset ? Number(request.params.arguments.offset) : 0;

          try {
            const result = await db.searchExploits(query, {
              platform,
              type,
              cve,
              author,
              startDate,
              endDate,
              verified,
              limit,
              offset
            });

            return {
              content: [{
                type: 'text',
                text: JSON.stringify({
                  total: result.total,
                  offset,
                  limit,
                  exploits: result.exploits
                }, null, 2)
              }]
            };
          } catch (error) {
            if (error instanceof McpError) {
              throw error;
            }
            throw new McpError(
              ErrorCode.InternalError,
              `Error searching exploits: ${(error as Error).message}`
            );
          }
        }

        case 'get_exploit': {
          const id = Number(request.params.arguments?.id);
          if (isNaN(id)) {
            throw new McpError(
              ErrorCode.InvalidParams,
              'Exploit ID must be a number'
            );
          }

          const includeCode = request.params.arguments?.include_code !== false;

          try {
            const exploit = await db.getExploitById(id);
            if (!exploit) {
              throw new McpError(
                ErrorCode.InvalidParams,
                `Exploit with ID ${id} not found`
              );
            }

            const result: any = { ...exploit };

            if (includeCode) {
              const code = await exploitdb.getExploitCode(exploit.file);
              if (code) {
                result.code = code;
              }
            }

            return {
              content: [{
                type: 'text',
                text: JSON.stringify(result, null, 2)
              }]
            };
          } catch (error) {
            if (error instanceof McpError) {
              throw error;
            }
            throw new McpError(
              ErrorCode.InternalError,
              `Error getting exploit: ${(error as Error).message}`
            );
          }
        }

        case 'find_by_cve': {
          const cve = String(request.params.arguments?.cve);
          if (!cve) {
            throw new McpError(
              ErrorCode.InvalidParams,
              'CVE ID is required'
            );
          }

          const limit = request.params.arguments?.limit ? Number(request.params.arguments.limit) : config.maxResults;

          try {
            const exploits = await db.findExploitsByCve(cve, limit);

            return {
              content: [{
                type: 'text',
                text: JSON.stringify({
                  cve,
                  count: exploits.length,
                  exploits
                }, null, 2)
              }]
            };
          } catch (error) {
            if (error instanceof McpError) {
              throw error;
            }
            throw new McpError(
              ErrorCode.InternalError,
              `Error finding exploits by CVE: ${(error as Error).message}`
            );
          }
        }

        case 'get_recent_exploits': {
          const limit = request.params.arguments?.limit ? Number(request.params.arguments.limit) : config.maxResults;

          try {
            const exploits = await db.getRecentExploits(limit);

            return {
              content: [{
                type: 'text',
                text: JSON.stringify({
                  count: exploits.length,
                  exploits
                }, null, 2)
              }]
            };
          } catch (error) {
            if (error instanceof McpError) {
              throw error;
            }
            throw new McpError(
              ErrorCode.InternalError,
              `Error getting recent exploits: ${(error as Error).message}`
            );
          }
        }

        case 'get_statistics': {
          try {
            const statistics = await db.getStatistics();

            return {
              content: [{
                type: 'text',
                text: JSON.stringify(statistics, null, 2)
              }]
            };
          } catch (error) {
            if (error instanceof McpError) {
              throw error;
            }
            throw new McpError(
              ErrorCode.InternalError,
              `Error getting statistics: ${(error as Error).message}`
            );
          }
        }

        default:
          throw new McpError(
            ErrorCode.MethodNotFound,
            `Unknown tool: ${request.params.name}`
          );
      }
    });
  }

  /**
   * Initialize the server
   */
  async init() {
    try {
      // Ensure data directory exists
      await ensureDataDir();

      // Initialize the database
      await db.initDatabase();

      // Set up automatic updates if configured
      if (config.updateInterval > 0) {
        this.updateInterval = setInterval(async () => {
          try {
            console.log('Running scheduled database update...');
            await exploitdb.updateDatabase();
            console.log('Scheduled database update completed');
          } catch (error) {
            console.error('Error in scheduled database update:', error);
          }
        }, config.updateInterval * 60 * 60 * 1000); // Convert hours to milliseconds
      }

      console.log('ExploitDB MCP server initialized');
    } catch (error) {
      console.error('Error initializing server:', error);
      throw error;
    }
  }

  /**
   * Start the server
   */
  async run() {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
    console.error('ExploitDB MCP server running on stdio');
  }

  /**
   * Close the server
   */
  async close() {
    if (this.updateInterval) {
      clearInterval(this.updateInterval);
      this.updateInterval = null;
    }

    await db.closeDatabase();
    await this.server.close();
  }
}

// Main function
const main = async () => {
  try {
    const server = new ExploitDBServer();
    await server.init();
    await server.run();
  } catch (error) {
    console.error('Server error:', error);
    process.exit(1);
  }
};

// Run the main function
main();