1
# frozen_string_literal: true
2

3
require 'msf/core/mcp'
4
require 'optparse'
5

6
module Msf::MCP
7
  # Main application class that orchestrates the MCP server startup and lifecycle
8
  class Application
9
    VERSION = '0.1.0'
10
    BANNER = <<~BANNER
11
      MSF MCP Server v#{VERSION}
12
      Model Context Protocol server for Metasploit Framework
13
    BANNER
14

15
    # For testing purposes:
16
    attr_reader :config, :msf_client, :mcp_server, :rate_limiter, :options, :rpc_manager
17

18
    # Initialize the application with command-line arguments
19
    #
20
    # @param argv [Array<String>] Command-line arguments
21
    # @param output [IO] Output stream for messages (default: $stderr)
22
    def initialize(argv = ARGV, output: $stderr)
23
      @argv = argv.dup
24
      @output = output
25
      @options = {}
26
      @config = nil
27
      @msf_client = nil
28
      @mcp_server = nil
29
      @rate_limiter = nil
30
      @rpc_manager = nil
31
    end
32

33
    # Run the application
34
    #
35
    # @return [void]
36
    def run
37
      parse_arguments
38
      install_signal_handlers
39
      load_configuration
40
      validate_configuration
41
      initialize_logger
42
      initialize_rate_limiter
43
      ensure_rpc_server
44
      initialize_metasploit_client
45
      authenticate_metasploit
46
      initialize_mcp_server
47
      start_mcp_server
48
    rescue Msf::MCP::Config::ValidationError, Msf::MCP::Config::ConfigurationError => e
49
      handle_configuration_error(e)
50
    rescue Msf::MCP::Metasploit::ConnectionError => e
51
      handle_connection_error(e)
52
    rescue Msf::MCP::Metasploit::APIError => e
53
      handle_api_error(e)
54
    rescue Msf::MCP::Metasploit::AuthenticationError => e
55
      handle_authentication_error(e)
56
    rescue Msf::MCP::Metasploit::RpcStartupError => e
57
      handle_rpc_startup_error(e)
58
    rescue StandardError => e
59
      handle_fatal_error(e)
60
    end
61

62
    # Shutdown the application gracefully
63
    #
64
    # Performs cleanup operations before process termination:
65
    # - Logs shutdown event via Rex
66
    # - Closes MCP server and Metasploit client connections
67
    # - Cleans up resources
68
    #
69
    # @param signal [String] Signal name (e.g., 'INT', 'TERM')
70
    # @return [void]
71
    def shutdown(signal = 'INT')
72
      ilog({
73
        message: 'Shutting down',
74
        context: { signal: "SIG#{signal}" }
75
      }, LOG_SOURCE, LOG_INFO)
76
      @mcp_server&.shutdown
77
      @rpc_manager&.stop_rpc_server
78
      @output.puts "\nShutdown complete"
79
    end
80

81
    private
82

83
    # Parse command-line arguments
84
    #
85
    # @return [void]
86
    def parse_arguments
87
      parser = OptionParser.new do |opts|
88
        opts.banner = BANNER + "\nUsage: msfmcp [options]"
89

90
        opts.on('--config PATH', 'Path to configuration file') do |path|
91
          @options[:config_path] = File.expand_path(path)
92
        end
93

94
        opts.on('--enable-logging', 'Enable file logging') do
95
          @options[:enable_logging_cli] = true
96
        end
97

98
        opts.on('--log-file PATH', 'Log file path (overrides config file)') do |path|
99
          @options[:log_file_cli] = path
100
        end
101

102
        opts.on('--user USER', 'MSF API username (for MessagePack auth)') do |user|
103
          @options[:msf_user_cli] = user
104
        end
105

106
        opts.on('--password PASS', 'MSF API password (for MessagePack auth)') do |password|
107
          @options[:msf_password_cli] = password
108
        end
109

110
        opts.on('--no-auto-start-rpc', 'Disable automatic RPC server startup') do
111
          @options[:no_auto_start_rpc] = true
112
        end
113

114
        opts.on('--mcp-transport TRANSPORT', 'MCP server transport type (\'stdio\' or \'http\')') do |transport|
115
          @options[:mcp_transport] = transport
116
        end
117

118
        opts.on('-h', '--help', 'Show this help message') do
119
          @output.puts opts
120
          exit 0
121
        end
122

123
        opts.on('-v', '--version', 'Show version information') do
124
          @output.puts "msfmcp version #{VERSION}"
125
          exit 0
126
        end
127
      end
128

129
      parser.parse!(@argv)
130
    end
131

132
    # Register a Rex log source when logging is enabled.
133
    #
134
    # Selects a JsonFlatfile sink pointed at the configured log path and wraps it
135
    # with the sanitizing middleware unless sanitization has been explicitly
136
    # disabled in the config.
137
    #
138
    # Priority: CLI flags > config file > defaults
139
    #
140
    # @return [void]
141
    def initialize_logger
142
      return unless @options[:enable_logging_cli] || @config.dig(:logging, :enabled)
143

144
      log_file = @options[:log_file_cli] || @config.dig(:logging, :log_file)
145
      level = @config.dig(:logging, :level)
146
      threshold = case @config.dig(:logging, :level).upcase
147
                  when 'DEBUG'
148
                    Rex::Logging::LEV_3
149
                  when 'INFO'
150
                    Rex::Logging::LEV_2
151
                  when 'WARN'
152
                    Rex::Logging::LEV_1
153
                  when 'ERROR'
154
                    Rex::Logging::LEV_0
155
                  end
156
      inner = Msf::MCP::Logging::Sinks::JsonFlatfile.new(log_file)
157
      sink  = @config.dig(:logging, :sanitize) ? Msf::MCP::Logging::Sinks::Sanitizing.new(inner) : inner
158

159
      deregister_log_source(LOG_SOURCE) if log_source_registered?(LOG_SOURCE)
160
      register_log_source(LOG_SOURCE, sink, threshold)
161
    end
162

163
    # Install signal handlers for graceful shutdown
164
    #
165
    # @return [void]
166
    def install_signal_handlers
167
      Signal.trap('INT') { shutdown('INT'); exit 0 }
168
      Signal.trap('TERM') { shutdown('TERM'); exit 0 }
169
    end
170

171
    # Load configuration from file or use defaults
172
    #
173
    # @return [void]
174
    def load_configuration
175
      if @options[:config_path]
176
        @output.puts "Loading configuration from #{@options[:config_path]}"
177
        @config = Msf::MCP::Config::Loader.load(@options[:config_path])
178
      else
179
        @output.puts "No configuration file specified, using defaults"
180
        @config = Msf::MCP::Config::Loader.load_from_hash({})
181
      end
182

183
      # Apply CLI authentication overrides (highest priority)
184
      if @options[:msf_user_cli]
185
        @config[:msf_api][:user] = @options[:msf_user_cli]
186
      end
187
      if @options[:msf_password_cli]
188
        @config[:msf_api][:password] = @options[:msf_password_cli]
189
      end
190
      if @options[:no_auto_start_rpc]
191
        @config[:msf_api][:auto_start_rpc] = false
192
      end
193
      if @options[:mcp_transport]
194
        @config[:mcp][:transport] = @options[:mcp_transport]
195
      end
196
    end
197

198
    # Validate the loaded configuration
199
    #
200
    # @return [void]
201
    def validate_configuration
202
      @output.puts "Validating configuration..."
203
      Msf::MCP::Config::Validator.validate!(@config)
204
      @output.puts "Configuration valid"
205
    end
206

207
    # Initialize the rate limiter
208
    #
209
    # @return [void]
210
    def initialize_rate_limiter
211
      @rate_limiter = Msf::MCP::Security::RateLimiter.new(
212
        requests_per_minute: @config.dig(:rate_limit, :requests_per_minute) || 60,
213
        burst_size: @config.dig(:rate_limit, :burst_size)
214
      )
215
    end
216

217
    # Ensure the Metasploit RPC server is available, auto-starting if needed
218
    #
219
    # @return [void]
220
    def ensure_rpc_server
221
      @rpc_manager = Msf::MCP::RpcManager.new(
222
        config: @config,
223
        output: @output
224
      )
225
      @rpc_manager.ensure_rpc_available
226
    end
227

228
    # Initialize the Metasploit client
229
    #
230
    # @return [void]
231
    def initialize_metasploit_client
232
      @output.puts "Connecting to Metasploit RPC at #{@config[:msf_api][:host]}:#{@config[:msf_api][:port]}"
233
      @msf_client = Msf::MCP::Metasploit::Client.new(
234
        api_type: @config[:msf_api][:type],
235
        host: @config[:msf_api][:host],
236
        port: @config[:msf_api][:port],
237
        endpoint: @config[:msf_api][:endpoint],
238
        token: @config[:msf_api][:token],
239
        ssl: @config[:msf_api][:ssl]
240
      )
241
    end
242

243
    # Authenticate with Metasploit if using MessagePack
244
    #
245
    # @return [void]
246
    def authenticate_metasploit
247
      if @config[:msf_api][:type] == 'messagepack'
248
        @output.puts "Authenticating with Metasploit..."
249
        @msf_client.authenticate(@config[:msf_api][:user].to_s, @config[:msf_api][:password].to_s)
250
        @output.puts "Authentication successful"
251
      else
252
        @output.puts "Using JSON-RPC with token authentication"
253
      end
254
    end
255

256
    # Initialize the MCP server
257
    #
258
    # @return [void]
259
    def initialize_mcp_server
260
      @output.puts "Initializing MCP server..."
261
      @mcp_server = Msf::MCP::Server.new(
262
        msf_client: @msf_client,
263
        rate_limiter: @rate_limiter
264
      )
265
    end
266

267
    # Start the MCP server with configured transport
268
    #
269
    # @return [void]
270
    def start_mcp_server
271
      transport = (@config.dig(:mcp, :transport) || 'stdio').to_sym
272
      host = @config.dig(:mcp, :host) || 'localhost'
273
      port = @config.dig(:mcp, :port) || 3000
274

275
      if transport == :http
276
        @output.puts "Starting MCP server on HTTP transport..."
277
        @output.puts "Server listening on http://#{host}:#{port}"
278
        @output.puts "Press Ctrl+C to shutdown"
279
        @mcp_server.start(transport: :http, host: host, port: port)
280
      else
281
        @output.puts "Starting MCP server on stdio transport..."
282
        @output.puts "Server ready - waiting for MCP requests"
283
        @output.puts "Press Ctrl+C to shutdown"
284
        @mcp_server.start(transport: :stdio)
285
      end
286
    end
287

288
    # Error handlers
289

290
    def handle_configuration_error(error)
291
      @output.puts "Configuration validation failed: #{error.message}"
292
      exit 1
293
    end
294

295
    def handle_connection_error(error)
296
      elog({
297
        message: 'Connection error',
298
        context: { host: @config[:msf_api][:host], port: @config[:msf_api][:port] },
299
        exception: error
300
      }, LOG_SOURCE, LOG_ERROR)
301
      @output.puts "Connection error to Metasploit RPC at #{@config[:msf_api][:host]}:#{@config[:msf_api][:port]} - #{error.message}"
302
      exit 1
303
    end
304

305
    def handle_api_error(error)
306
      elog({ message: 'Metasploit API error', exception: error }, LOG_SOURCE, LOG_ERROR)
307
      @output.puts "Metasploit API error: #{error.message}"
308
      exit 1
309
    end
310

311
    def handle_authentication_error(error)
312
      elog({
313
        message: 'Authentication error',
314
        context: { username: @config[:msf_api][:user].to_s },
315
        exception: error
316
      }, LOG_SOURCE, LOG_ERROR)
317
      @output.puts "Authentication error (username: #{@config[:msf_api][:user]}): #{error.message}"
318
      exit 1
319
    end
320

321
    def handle_rpc_startup_error(error)
322
      elog({ message: 'RPC startup error', exception: error }, LOG_SOURCE, LOG_ERROR)
323
      @output.puts "RPC startup error: #{error.message}"
324
      exit 1
325
    end
326

327
    def handle_fatal_error(error)
328
      elog({ message: 'Fatal error during startup', exception: error }, LOG_SOURCE, LOG_ERROR)
329
      @output.puts "Fatal error: #{error.message}"
330
      @output.puts error.backtrace.first(5).join("\n") if error.backtrace
331
      exit 1
332
    end
333
  end
334
end
335

336