1
# frozen_string_literal: true
2

3
require 'spec_helper'
4
require 'rex/text'
5
require Metasploit::Framework.root.join('plugins/mcp.rb').to_path
6

7
RSpec.describe Msf::Plugin::MCP do
8
  include_context 'Msf::UIDriver'
9

10
  let(:framework) { instance_double(Msf::Framework) }
11
  let(:framework_datastore) { { 'VERBOSE' => false } }
12
  let(:output) { driver_output }
13
  let(:base_opts) { { 'LocalOutput' => output } }
14

15
  let(:mock_thread) do
16
    instance_double(Thread, alive?: false, join: true, kill: nil)
17
  end
18

19
  let(:threads_manager) do
20
    instance_double('Msf::Framework::ThreadManager').tap do |tm|
21
      allow(tm).to receive(:spawn).and_return(mock_thread)
22
    end
23
  end
24

25
  let(:plugins_collection) do
26
    instance_double('Msf::PluginManager').tap do |pm|
27
      allow(pm).to receive(:find).and_return(nil)
28
      allow(pm).to receive(:load).and_return(true)
29
      allow(pm).to receive(:unload).and_return(true)
30
    end
31
  end
32

33
  let(:mock_client_class) do
34
    Class.new do
35
      def initialize(**_args); end
36

37
      def authenticate(*_args)
38
        'token'
39
      end
40
    end
41
  end
42

43
  let(:mock_rate_limiter_class) do
44
    Class.new do
45
      def initialize(**_args); end
46
    end
47
  end
48

49
  let(:mock_server_class) do
50
    Class.new do
51
      def initialize(**_args); end
52
      def start(**_args); end
53
      def shutdown; end
54
    end
55
  end
56

57
  before do
58
    allow(framework).to receive(:threads).and_return(threads_manager)
59
    allow(framework).to receive(:plugins).and_return(plugins_collection)
60
    allow(framework).to receive(:datastore).and_return(framework_datastore)
61

62
    stub_const('Msf::MCP::Metasploit::Client', mock_client_class)
63
    stub_const('Msf::MCP::Metasploit::AuthenticationError', Class.new(StandardError))
64
    stub_const('Msf::MCP::Metasploit::ConnectionError', Class.new(StandardError))
65
    stub_const('Msf::MCP::Security::RateLimiter', mock_rate_limiter_class)
66
    stub_const('Msf::MCP::Server', mock_server_class)
67

68
    allow(Rex::Text).to receive(:rand_text_alphanumeric).with(12).and_return('abcdefghijkl')
69
    allow_any_instance_of(described_class).to receive(:sleep)
70

71
    mock_dispatcher = instance_double(Msf::Plugin::MCP::McpCommandDispatcher)
72
    allow(mock_dispatcher).to receive(:plugin=)
73
    allow_any_instance_of(Msf::Plugin::MCP).to receive(:add_console_dispatcher).and_return(mock_dispatcher)
74
    allow_any_instance_of(Msf::Plugin::MCP).to receive(:remove_console_dispatcher)
75
  end
76

77
  describe '#initialize' do
78
    subject(:plugin) { described_class.new(framework, base_opts) }
79

80
    it 'creates the plugin successfully without starting the server' do
81
      expect { plugin }.not_to raise_error
82
    end
83

84
    it 'does not create an MCP client' do
85
      expect(mock_client_class).not_to receive(:new)
86
      plugin
87
    end
88

89
    it 'does not spawn a server thread' do
90
      expect(threads_manager).not_to receive(:spawn)
91
      plugin
92
    end
93

94
    it 'prints a message about using mcp start' do
95
      plugin
96
      expect(@output.join("\n")).to include('mcp start')
97
    end
98

99
    it 'registers the command dispatcher' do
100
      expect_any_instance_of(described_class).to receive(:add_console_dispatcher)
101
        .with(described_class::McpCommandDispatcher)
102
        .and_return(instance_double(described_class::McpCommandDispatcher, :plugin= => nil))
103
      plugin
104
    end
105

106
    it 'has nil server_config' do
107
      expect(plugin.server_config).to be_nil
108
    end
109
  end
110

111
  describe '#start_server' do
112
    subject(:plugin) { described_class.new(framework, base_opts) }
113

114
    context 'with default options' do
115
      before { plugin.start_server({}) }
116

117
      it 'creates an MCP client' do
118
        expect(plugin.msf_client).not_to be_nil
119
      end
120

121
      it 'starts the server in a background thread' do
122
        expect(threads_manager).to have_received(:spawn).with('MCPServer', false)
123
      end
124

125
      it 'prints a status message with the listening address' do
126
        expect(@output.join("\n")).to include('MCP server started on localhost:3000 (transport: http)')
127
      end
128

129
      it 'stores the server configuration' do
130
        expect(plugin.server_config).to be_a(Hash)
131
        expect(plugin.server_config[:mcp][:transport]).to eq('http')
132
      end
133
    end
134

135
    context 'with custom options' do
136
      before { plugin.start_server('ServerHost' => '0.0.0.0', 'ServerPort' => '8080') }
137

138
      it 'applies the provided host and port' do
139
        expect(plugin.server_config[:mcp][:host]).to eq('0.0.0.0')
140
        expect(plugin.server_config[:mcp][:port]).to eq(8080)
141
      end
142
    end
143

144
    context 'with stdio transport' do
145
      it 'rejects stdio as an unknown option since Transport is no longer accepted' do
146
        dispatcher = Msf::Plugin::MCP::McpCommandDispatcher.new(driver)
147
        dispatcher.plugin = plugin
148
        capture_logging(dispatcher)
149
        dispatcher.cmd_mcp('start', 'Transport=stdio')
150
        expect(@error.join("\n")).to include('Unknown option: Transport')
151
      end
152
    end
153

154
    context 'when server is already running' do
155
      before do
156
        plugin.start_server({})
157
        reset_logging!
158
      end
159

160
      it 'prints an error and does not restart' do
161
        plugin.start_server({})
162
        expect(@error.join("\n")).to include('MCP server is already running')
163
      end
164
    end
165

166
    context 'when port is already in use' do
167
      let(:failing_server_class) do
168
        Class.new do
169
          def initialize(**_args); end
170

171
          def start(**_args)
172
            raise Errno::EADDRINUSE
173
          end
174

175
          def shutdown; end
176
        end
177
      end
178

179
      before do
180
        stub_const('Msf::MCP::Server', failing_server_class)
181
        allow(threads_manager).to receive(:spawn) do |_name, _critical, &block|
182
          block.call
183
        end
184
      end
185

186
      it 'prints an error with address-in-use message' do
187
        plugin.start_server({})
188
        expect(@error.join("\n")).to include('Address already in use')
189
      end
190
    end
191

192
    context 'when RPC authentication fails' do
193
      let(:failing_client_class) do
194
        error_class = Msf::MCP::Metasploit::AuthenticationError
195
        Class.new do
196
          define_method(:initialize) { |**_args| }
197
          define_method(:authenticate) { |*_args| raise error_class, 'bad credentials' }
198
        end
199
      end
200

201
      before do
202
        stub_const('Msf::MCP::Metasploit::Client', failing_client_class)
203
      end
204

205
      it 'prints an error with authentication failure message' do
206
        plugin.start_server({})
207
        expect(@error.join("\n")).to include('RPC authentication failed')
208
      end
209
    end
210

211
    context 'when RPC connection fails' do
212
      let(:failing_client_class) do
213
        error_class = Msf::MCP::Metasploit::ConnectionError
214
        Class.new do
215
          define_method(:initialize) { |**_args| }
216
          define_method(:authenticate) { |*_args| raise error_class, 'connection refused' }
217
        end
218
      end
219

220
      before do
221
        stub_const('Msf::MCP::Metasploit::Client', failing_client_class)
222
      end
223

224
      it 'prints an error with connection failure message' do
225
        plugin.start_server({})
226
        expect(@error.join("\n")).to include('RPC connection failed')
227
      end
228

229
      it 'unloads the auto-started msgrpc plugin on failure' do
230
        msgrpc_plugin = instance_double('Msf::Plugin::MSGRPC', name: 'msgrpc')
231
        call_count = 0
232
        allow(plugins_collection).to receive(:find) do
233
          call_count += 1
234
          call_count > 1 ? msgrpc_plugin : nil
235
        end
236
        expect(plugins_collection).to receive(:unload).with(msgrpc_plugin)
237
        plugin.start_server({})
238
      end
239
    end
240
  end
241

242
  describe '#stop_server' do
243
    subject(:plugin) { described_class.new(framework, base_opts) }
244

245
    context 'when server is running' do
246
      before do
247
        plugin.start_server({})
248
        reset_logging!
249
      end
250

251
      it 'stops the server and prints a message' do
252
        plugin.stop_server
253
        expect(@output.join("\n")).to include('MCP server stopped')
254
      end
255

256
      it 'nils out the server reference' do
257
        plugin.stop_server
258
        expect(plugin.mcp_server).to be_nil
259
      end
260

261
      it 'nils out the client reference' do
262
        plugin.stop_server
263
        expect(plugin.msf_client).to be_nil
264
      end
265
    end
266

267
    context 'when server is not running' do
268
      it 'prints an error' do
269
        plugin.stop_server
270
        expect(@error.join("\n")).to include('MCP server is already stopped')
271
      end
272
    end
273
  end
274

275
  describe '#restart_server' do
276
    subject(:plugin) { described_class.new(framework, base_opts) }
277

278
    context 'when server is running' do
279
      before do
280
        plugin.start_server({})
281
        reset_logging!
282
      end
283

284
      it 'restarts with new options' do
285
        plugin.restart_server('ServerPort' => '9090')
286
        expect(plugin.server_config[:mcp][:port]).to eq(9090)
287
      end
288
    end
289

290
    context 'when server is not running' do
291
      it 'starts the server fresh' do
292
        plugin.restart_server('ServerPort' => '8080')
293
        expect(plugin.server_config[:mcp][:port]).to eq(8080)
294
        expect(plugin.mcp_server).not_to be_nil
295
      end
296
    end
297
  end
298

299
  describe '#print_mcp_status' do
300
    subject(:plugin) { described_class.new(framework, base_opts) }
301

302
    context 'when server has never been configured' do
303
      it 'prints not configured status' do
304
        plugin.print_mcp_status
305
        expect(@output.join("\n")).to include('stopped (not configured)')
306
      end
307
    end
308

309
    context 'when server is running' do
310
      before do
311
        plugin.start_server({})
312
        reset_logging!
313
      end
314

315
      it 'prints running status with details' do
316
        plugin.print_mcp_status
317
        combined = @output.join("\n")
318
        expect(combined).to include('MCP server status: running')
319
        expect(combined).to include('http://localhost:3000')
320
      end
321
    end
322

323
    context 'when server was started then stopped' do
324
      before do
325
        plugin.start_server({})
326
        plugin.stop_server
327
        reset_logging!
328
      end
329

330
      it 'prints stopped status with last known config' do
331
        plugin.print_mcp_status
332
        combined = @output.join("\n")
333
        expect(combined).to include('MCP server status: stopped')
334
      end
335
    end
336
  end
337

338
  describe '#cleanup' do
339
    subject(:plugin) { described_class.new(framework, base_opts) }
340

341
    before do
342
      allow(plugin).to receive(:remove_console_dispatcher)
343
    end
344

345
    context 'when server is running' do
346
      before do
347
        plugin.start_server({})
348
        reset_logging!
349
      end
350

351
      it 'shuts down the MCP server' do
352
        server = plugin.mcp_server
353
        expect(server).to receive(:shutdown)
354
        plugin.cleanup
355
      end
356

357
      it 'prints a stop message' do
358
        plugin.cleanup
359
        expect(@output.join("\n")).to include('MCP server stopped')
360
      end
361

362
      it 'nils out all references' do
363
        plugin.cleanup
364
        expect(plugin.mcp_server).to be_nil
365
        expect(plugin.server_thread).to be_nil
366
        expect(plugin.msf_client).to be_nil
367
        expect(plugin.rate_limiter).to be_nil
368
        expect(plugin.started_at).to be_nil
369
      end
370
    end
371

372
    context 'when server was never started' do
373
      it 'deregisters the console dispatcher without error' do
374
        expect(plugin).to receive(:remove_console_dispatcher).with('MCP')
375
        expect { plugin.cleanup }.not_to raise_error
376
      end
377
    end
378

379
    context 'when msgrpc was auto-started' do
380
      before { plugin.start_server({}) }
381

382
      it 'unloads the auto-started msgrpc plugin' do
383
        plugin.auto_started_rpc = true
384
        msgrpc_plugin = instance_double('Msf::Plugin::MSGRPC', name: 'msgrpc')
385
        allow(plugins_collection).to receive(:find).and_return(msgrpc_plugin)
386
        expect(plugins_collection).to receive(:unload).with(msgrpc_plugin)
387
        plugin.cleanup
388
      end
389
    end
390

391
    context 'when msgrpc was pre-existing' do
392
      before { plugin.start_server({}) }
393

394
      it 'does not unload the msgrpc plugin' do
395
        plugin.auto_started_rpc = false
396
        expect(plugins_collection).not_to receive(:unload)
397
        plugin.cleanup
398
      end
399
    end
400

401
    context 'when msgrpc unload fails' do
402
      before { plugin.start_server({}) }
403

404
      it 'prints a warning and continues cleanup' do
405
        plugin.auto_started_rpc = true
406
        msgrpc_plugin = instance_double('Msf::Plugin::MSGRPC', name: 'msgrpc')
407
        allow(plugins_collection).to receive(:find).and_return(msgrpc_plugin)
408
        allow(plugins_collection).to receive(:unload).and_raise(StandardError, 'unload failed')
409

410
        plugin.cleanup
411

412
        expect(@error.join("\n")).to include('Failed to unload auto-started msgrpc')
413
        expect(plugin.mcp_server).to be_nil
414
      end
415
    end
416
  end
417

418
  describe '#terminate_server_thread' do
419
    subject(:plugin) { described_class.new(framework, base_opts) }
420

421
    context 'when thread terminates within 5 seconds' do
422
      it 'does not force kill the thread' do
423
        alive_thread = instance_double(Thread, alive?: true, join: true)
424
        plugin.server_thread = alive_thread
425
        expect(alive_thread).not_to receive(:kill)
426
        plugin.send(:terminate_server_thread)
427
      end
428
    end
429

430
    context 'when thread does not terminate within 5 seconds' do
431
      it 'force kills the thread' do
432
        stuck_thread = instance_double(Thread, alive?: true, join: nil, kill: nil)
433
        plugin.server_thread = stuck_thread
434
        expect(stuck_thread).to receive(:kill)
435
        plugin.send(:terminate_server_thread)
436
      end
437

438
      it 'prints a warning message' do
439
        stuck_thread = instance_double(Thread, alive?: true, join: nil, kill: nil)
440
        plugin.server_thread = stuck_thread
441
        plugin.send(:terminate_server_thread)
442
        expect(@error.join("\n")).to include('did not terminate gracefully')
443
      end
444
    end
445

446
    context 'when thread is already dead' do
447
      it 'does nothing' do
448
        dead_thread = instance_double(Thread, alive?: false)
449
        plugin.server_thread = dead_thread
450
        expect(dead_thread).not_to receive(:join)
451
        expect(dead_thread).not_to receive(:kill)
452
        plugin.send(:terminate_server_thread)
453
      end
454
    end
455

456
    context 'when server_thread is nil' do
457
      it 'does nothing' do
458
        plugin.server_thread = nil
459
        expect { plugin.send(:terminate_server_thread) }.not_to raise_error
460
      end
461
    end
462
  end
463

464
  describe 'RPC resolution' do
465
    subject(:plugin) { described_class.new(framework, base_opts) }
466

467
    context 'with explicit RPC credentials' do
468
      it 'uses the provided credentials' do
469
        plugin.start_server('RpcUser' => 'admin', 'RpcPass' => 'secret123')
470
        expect(plugin.server_config[:rpc][:user]).to eq('admin')
471
        expect(plugin.server_config[:rpc][:pass]).to eq('secret123')
472
      end
473
    end
474

475
    context 'when msgrpc is already loaded (introspection path)' do
476
      let(:mock_msgrpc_server) do
477
        instance_double('Msf::RPC::Service').tap do |s|
478
          allow(s).to receive(:users).and_return({ 'admin' => 'secretpass' })
479
          allow(s).to receive(:srvhost).and_return('127.0.0.1')
480
          allow(s).to receive(:srvport).and_return(55_553)
481
          allow(s).to receive(:options).and_return({ ssl: true })
482
        end
483
      end
484

485
      let(:mock_msgrpc_plugin) do
486
        instance_double('Msf::Plugin::MSGRPC', name: 'msgrpc', server: mock_msgrpc_server)
487
      end
488

489
      before do
490
        allow(plugins_collection).to receive(:find) do |&block|
491
          [mock_msgrpc_plugin].find(&block)
492
        end
493
      end
494

495
      it 'uses credentials from the loaded msgrpc plugin' do
496
        plugin.start_server({})
497
        expect(plugin.server_config[:rpc][:user]).to eq('admin')
498
        expect(plugin.server_config[:rpc][:pass]).to eq('secretpass')
499
      end
500

501
      it 'uses host and port from the loaded msgrpc plugin' do
502
        plugin.start_server({})
503
        expect(plugin.server_config[:rpc][:host]).to eq('127.0.0.1')
504
        expect(plugin.server_config[:rpc][:port]).to eq(55_553)
505
      end
506

507
      it 'does not auto-start msgrpc' do
508
        expect(plugins_collection).not_to receive(:load)
509
        plugin.start_server({})
510
      end
511

512
      it 'does not set auto_started_rpc flag' do
513
        plugin.start_server({})
514
        expect(plugin.auto_started_rpc).to be false
515
      end
516
    end
517

518
    context 'when no msgrpc is loaded and no creds provided (auto-start path)' do
519
      it 'auto-starts msgrpc and prints credentials' do
520
        expect(plugins_collection).to receive(:load).with('msgrpc', hash_including('Pass' => 'abcdefghijkl', 'User' => 'msf'))
521
        plugin.start_server({})
522
        expect(@output.join("\n")).to include('Auto-started msgrpc')
523
        expect(@output.join("\n")).to include('abcdefghijkl')
524
      end
525

526
      it 'sets auto_started_rpc flag' do
527
        plugin.start_server({})
528
        expect(plugin.auto_started_rpc).to be true
529
      end
530
    end
531
  end
532

533
  describe 'option validation' do
534
    subject(:plugin) { described_class.new(framework, base_opts) }
535

536
    context 'with invalid ServerPort' do
537
      it 'prints an error' do
538
        plugin.start_server('ServerPort' => '99999')
539
        expect(@error.join("\n")).to include('Invalid value for ServerPort')
540
      end
541
    end
542

543
    context 'with invalid Transport' do
544
      it 'rejects Transport as an unknown option' do
545
        dispatcher = Msf::Plugin::MCP::McpCommandDispatcher.new(driver)
546
        dispatcher.plugin = plugin
547
        capture_logging(dispatcher)
548
        dispatcher.cmd_mcp('start', 'Transport=websocket')
549
        expect(@error.join("\n")).to include('Unknown option: Transport')
550
      end
551
    end
552

553
    context 'with invalid RpcSSL' do
554
      it 'prints an error' do
555
        plugin.start_server('RpcSSL' => 'maybe')
556
        expect(@error.join("\n")).to include('Invalid value for RpcSSL')
557
      end
558
    end
559

560
    context 'with invalid RateLimit' do
561
      it 'prints an error' do
562
        plugin.start_server('RateLimit' => '0')
563
        expect(@error.join("\n")).to include('Invalid value for RateLimit')
564
      end
565
    end
566

567
    context 'with RpcUser but no RpcPass' do
568
      it 'prints an error' do
569
        plugin.start_server('RpcUser' => 'admin')
570
        expect(@error.join("\n")).to include('Invalid value for RpcPass')
571
      end
572
    end
573

574
    context 'with RpcPass but no RpcUser' do
575
      it 'prints an error' do
576
        plugin.start_server('RpcPass' => 'secret')
577
        expect(@error.join("\n")).to include('Invalid value for RpcUser')
578
      end
579
    end
580
  end
581

582
  describe Msf::Plugin::MCP::McpCommandDispatcher do
583
    let(:plugin_instance) { Msf::Plugin::MCP.new(framework, base_opts) }
584
    let(:dispatcher) do
585
      d = Msf::Plugin::MCP::McpCommandDispatcher.new(driver)
586
      d.plugin = plugin_instance
587
      d
588
    end
589

590
    before do
591
      capture_logging(dispatcher)
592
    end
593

594
    describe '#cmd_mcp with start subcommand' do
595
      it 'parses Key=Value options and passes them to start_server' do
596
        expect(plugin_instance).to receive(:start_server) do |opts|
597
          expect(opts).to eq('RateLimit' => '120')
598
        end
599
        dispatcher.cmd_mcp('start', 'RateLimit=120')
600
      end
601

602
      it 'starts with empty opts when no options given' do
603
        expect(plugin_instance).to receive(:start_server) do |opts|
604
          expect(opts).to eq({})
605
        end
606
        dispatcher.cmd_mcp('start')
607
      end
608

609
      it 'rejects malformed options' do
610
        expect(plugin_instance).not_to receive(:start_server)
611
        dispatcher.cmd_mcp('start', 'badoption')
612
        expect(@error.join("\n")).to include('Invalid option format')
613
      end
614

615
      it 'rejects unknown option keys' do
616
        expect(plugin_instance).not_to receive(:start_server)
617
        dispatcher.cmd_mcp('start', 'FakeOption=value')
618
        expect(@error.join("\n")).to include('Unknown option: FakeOption')
619
      end
620
    end
621

622
    describe '#cmd_mcp with restart subcommand' do
623
      it 'parses options and passes them to restart_server' do
624
        expect(plugin_instance).to receive(:restart_server) do |opts|
625
          expect(opts).to eq('ServerPort' => '9090')
626
        end
627
        dispatcher.cmd_mcp('restart', 'ServerPort=9090')
628
      end
629
    end
630

631
    describe '#cmd_mcp with help subcommand' do
632
      it 'prints usage information including option examples' do
633
        dispatcher.cmd_mcp('help')
634
        combined = @output.join("\n")
635
        expect(combined).to include('Usage: mcp <subcommand> [options]')
636
        expect(combined).to include('ServerHost=<host>')
637
        expect(combined).to include('mcp start RpcUser=msf RpcPass=secret')
638
      end
639
    end
640

641
    describe '#cmd_mcp_tabs' do
642
      it 'returns subcommands when typing subcommand' do
643
        # User typed: "mcp <tab>" → words = ['mcp'], str = ''
644
        result = dispatcher.cmd_mcp_tabs('', ['mcp'])
645
        expect(result).to contain_exactly('status', 'start', 'stop', 'restart', 'help')
646
      end
647

648
      it 'returns option completions for start subcommand' do
649
        # User typed: "mcp start <tab>" → words = ['mcp', 'start'], str = ''
650
        result = dispatcher.cmd_mcp_tabs('', ['mcp', 'start'])
651
        expect(result).to include('ServerHost=', 'ServerPort=', 'RpcHost=', 'RpcPass=')
652
      end
653

654
      it 'filters options by partial input' do
655
        # User typed: "mcp start Rpc<tab>" → words = ['mcp', 'start'], str = 'Rpc'
656
        result = dispatcher.cmd_mcp_tabs('Rpc', ['mcp', 'start'])
657
        expect(result).to contain_exactly('RpcHost=', 'RpcPort=', 'RpcUser=', 'RpcPass=', 'RpcSSL=')
658
      end
659

660
      it 'filters options case-insensitively' do
661
        # User typed: "mcp start rpc<tab>" → words = ['mcp', 'start'], str = 'rpc'
662
        result = dispatcher.cmd_mcp_tabs('rpc', ['mcp', 'start'])
663
        expect(result).to contain_exactly('RpcHost=', 'RpcPort=', 'RpcUser=', 'RpcPass=', 'RpcSSL=')
664
      end
665

666
      it 'returns empty for non-start/restart subcommands' do
667
        # User typed: "mcp status <tab>" → words = ['mcp', 'status'], str = ''
668
        result = dispatcher.cmd_mcp_tabs('', ['mcp', 'status'])
669
        expect(result).to eq([])
670
      end
671
    end
672
  end
673
end
674

675