1
# frozen_string_literal: true
2

3
module Msf::MCP
4
  ##
5
  # MCP Server Wrapper for Metasploit Framework
6
  #
7
  # This class initializes and manages the MCP server with all registered tools.
8
  # It provides a clean interface for starting/stopping the server and integrates
9
  # with the Metasploit client and security layers.
10
  #
11
  # The Server expects fully configured and authenticated dependencies to be
12
  # provided during initialization. It does not handle configuration loading
13
  # or client authentication - those are responsibilities of the calling code.
14
  #
15
  class Server
16

17
    ##
18
    # Initialize the MCP server with required dependencies
19
    #
20
    # @param msf_client [Metasploit::Client] Configured and authenticated Metasploit client
21
    # @param rate_limiter [Security::RateLimiter] Configured rate limiter
22
    #
23
    def initialize(msf_client:, rate_limiter:)
24
      @msf_client = msf_client
25

26
      # Create server context (passed to all tool calls)
27
      # Tools only need msf_client and rate_limiter
28
      @server_context = {
29
        msf_client: @msf_client,
30
        rate_limiter: rate_limiter
31
      }
32

33
      # Create MCP configuration with request lifecycle callbacks
34
      mcp_config = ::MCP::Configuration.new
35
      mcp_config.around_request = create_around_request
36
      mcp_config.exception_reporter = create_exception_reporter
37

38
      # Initialize MCP server with all tools
39
      @mcp_server = ::MCP::Server.new(
40
        name: 'msfmcp',
41
        version: Msf::MCP::Application::VERSION,
42
        tools: [
43
          Tools::SearchModules,
44
          Tools::ModuleInfo,
45
          Tools::HostInfo,
46
          Tools::ServiceInfo,
47
          Tools::VulnerabilityInfo,
48
          Tools::NoteInfo,
49
          Tools::CredentialInfo,
50
          Tools::LootInfo
51
        ],
52
        server_context: @server_context,
53
        configuration: mcp_config
54
      )
55
    end
56

57
    ##
58
    # Start the MCP server with specified transport
59
    #
60
    # @param transport [Symbol] Transport type (:stdio or :http)
61
    # @param host [String] Host address for HTTP transport (default: 'localhost')
62
    # @param port [Integer] Port number for HTTP transport (default: 3000)
63
    #
64
    # @return [MCP::Server] The MCP server instance (for testing purposes)
65
    # @raise [ArgumentError] If an unknown transport is specified
66
    #
67
    def start(transport: :stdio, host: 'localhost', port: 3000)
68
      case transport
69
      when :stdio
70
        start_stdio
71
      when :http
72
        start_http(host, port)
73
      else
74
        raise ArgumentError, "Unknown transport: #{transport}. Use :stdio or :http"
75
      end
76
    end
77

78
    ##
79
    # Shutdown the MCP server and cleanup resources
80
    #
81
    def shutdown
82
      @msf_client&.shutdown
83
      @mcp_server = nil
84
    end
85

86
    private
87

88
    ##
89
    # Start stdio transport (for CLI usage)
90
    #
91
    # @return [MCP::Server] The MCP server instance (for testing purposes)
92
    #
93
    def start_stdio
94
      transport = ::MCP::Server::Transports::StdioTransport.new(@mcp_server)
95
      transport.open
96
      @mcp_server
97
    end
98

99
    ##
100
    # Start HTTP transport (for web/network usage)
101
    #
102
    # The transport implements the Rack app interface (#call), so it is mounted
103
    # directly. MCP-aware request/response logging is handled by the
104
    # Middleware::RequestLogger middleware.
105
    #
106
    # @param host [String] Host address to bind to
107
    # @param port [Integer] Port to listen on
108
    #
109
    # @return [MCP::Server] The MCP server instance (for testing purposes)
110
    #
111
    def start_http(host, port)
112
      require 'rackup'
113
      require 'rack/handler/puma'
114

115
      transport = ::MCP::Server::Transports::StreamableHTTPTransport.new(@mcp_server)
116

117
      # Build the Rack application with logging middleware.
118
      # The transport itself is a Rack app (implements #call).
119
      rack_app = Rack::Builder.new do
120
        use Msf::MCP::Middleware::RequestLogger
121
        run transport
122
      end
123

124
      Rackup::Handler::Puma.run(
125
        rack_app,
126
        Port: port,
127
        Host: host,
128
        Silent: true
129
      )
130

131
      @mcp_server
132
    end
133

134
    ##
135
    # Create around_request callback for MCP SDK
136
    #
137
    # This callback wraps every JSON-RPC request handler, providing access to
138
    # both the instrumentation data and the response result. It replaces the
139
    # deprecated +instrumentation_callback+ which only fires after completion
140
    # and does not expose the result.
141
    #
142
    # The +data+ hash is populated by the SDK with:
143
    # - :method — the JSON-RPC method name (e.g. "tools/call", "tools/list")
144
    # - :tool_name, :prompt_name, :resource_uri — specific handler identifiers
145
    # - :tool_arguments — arguments passed to a tool call
146
    # - :client — client info hash (name, version)
147
    # - :error — error type symbol (e.g. :tool_not_found, :internal_error)
148
    # - :duration — added in the ensure block after this callback returns
149
    #
150
    # @return [Proc] Callback that wraps request execution and logs via Rex
151
    #
152
    def create_around_request
153
      ->(data, &request_handler) do
154
        result = request_handler.call
155

156
        # Build message based on the type of request
157
        message = if data[:error]
158
                    "MCP Error: #{data[:error]}"
159
                  elsif data[:tool_name]
160
                    "Tool call: #{data[:tool_name]}"
161
                  elsif data[:prompt_name]
162
                    "Prompt call: #{data[:prompt_name]}"
163
                  elsif data[:resource_uri]
164
                    "Resource call: #{data[:resource_uri]}"
165
                  elsif data[:method]
166
                    "Method call: #{data[:method]}"
167
                  else
168
                    "MCP request"
169
                  end
170

171
        context = data.dup
172
        if result
173
          message = "#{message} (ERROR)" if result[:isError]
174
          context[:result] = result
175
        end
176

177
        if data[:error] || result&.fetch(:isError, nil)
178
          elog({ message: message, context: context }, LOG_SOURCE, LOG_ERROR)
179
        else
180
          ilog({ message: message, context: context }, LOG_SOURCE, LOG_INFO)
181
        end
182

183
        result
184
      end
185
    end
186

187
    ##
188
    # Create exception reporter callback for MCP SDK
189
    #
190
    # This callback is invoked for any server exception during request processing,
191
    # which are not tool execution errors.
192
    # It receives:
193
    # - exception: The Ruby exception object
194
    # - context: Hash with :request (JSON string) or :notification (method name string)
195
    #
196
    # @return [Proc] Callback that logs exceptions via Rex
197
    #
198
    def create_exception_reporter
199
      ->(exception, context) do
200
        return unless exception || context
201

202
        # Determine the context type and parse data
203
        error_context = {}
204

205
        if context&.fetch(:request, nil)
206
          error_context[:type] = 'request'
207
          request = nil
208
          begin
209
            request = JSON.parse(context[:request])
210
          rescue JSON::ParserError
211
            # Not valid JSON, log raw data
212
            error_context[:raw_data] = context[:request].inspect
213
          else
214
            error_context[:method] = request['method'] if request['method']
215
            error_context[:params] = request['params'] if request['params']
216
          end
217
        elsif context&.fetch(:notification, nil)
218
          error_context[:type] = 'notification'
219
          # context[:notification] is the notification method name (string)
220
          error_context[:method] = context[:notification]
221
        else
222
          error_context[:type] = 'unknown'
223
          error_context[:raw_data] = context.inspect
224
        end
225

226
        elog({
227
          message: "Error during #{error_context[:type]} processing#{error_context[:method] ? " (#{error_context[:method]})" : ''}",
228
          exception: exception,
229
          context: error_context
230
        }, LOG_SOURCE, LOG_ERROR)
231
      end
232
    end
233
  end
234
end
235

236