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::McpCommandDispatcher do
8
  include_context 'Msf::UIDriver'
9

10
  let(:framework) { instance_double(Msf::Framework) }
11
  let(:output) { driver_output }
12

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

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

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

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

35
      def authenticate(*_args)
36
        'token'
37
      end
38
    end
39
  end
40

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

47
  let(:mock_server_class) do
48
    Class.new do
49
      def initialize(**_args); end
50

51
      def start(**_args); end
52

53
      def shutdown; end
54
    end
55
  end
56

57
  let(:plugin) do
58
    described_class.new(driver).tap do |d|
59
      d.plugin = mcp_plugin
60
    end
61
  end
62

63
  let(:base_opts) { { 'LocalOutput' => output } }
64

65
  let(:mcp_plugin) { Msf::Plugin::MCP.new(framework, base_opts) }
66

67
  before do
68
    allow(framework).to receive(:threads).and_return(threads_manager)
69
    allow(framework).to receive(:plugins).and_return(plugins_collection)
70

71
    stub_const('Msf::MCP::Metasploit::Client', mock_client_class)
72
    stub_const('Msf::MCP::Metasploit::AuthenticationError', Class.new(StandardError))
73
    stub_const('Msf::MCP::Metasploit::ConnectionError', Class.new(StandardError))
74
    stub_const('Msf::MCP::Security::RateLimiter', mock_rate_limiter_class)
75
    stub_const('Msf::MCP::Server', mock_server_class)
76

77
    allow(Rex::Text).to receive(:rand_text_alphanumeric).with(12).and_return('abcdefghijkl')
78

79
    mock_dispatcher = instance_double(described_class)
80
    allow(mock_dispatcher).to receive(:plugin=)
81
    allow_any_instance_of(Msf::Plugin::MCP).to receive(:add_console_dispatcher).and_return(mock_dispatcher)
82
    allow_any_instance_of(Msf::Plugin::MCP).to receive(:remove_console_dispatcher)
83

84
    # Capture output from the plugin's print methods
85
    capture_logging(mcp_plugin)
86
  end
87

88
  describe '#name' do
89
    it 'returns MCP' do
90
      expect(plugin.name).to eq('MCP')
91
    end
92
  end
93

94
  describe '#commands' do
95
    it 'registers the mcp command' do
96
      expect(plugin.commands).to eq({ 'mcp' => 'Manage the MCP server' })
97
    end
98
  end
99

100
  describe '#cmd_mcp_tabs' do
101
    it 'returns all subcommands when typing subcommand' do
102
      # User typed: "mcp <tab>" → words = ['mcp'], str = ''
103
      expect(plugin.cmd_mcp_tabs('', ['mcp'])).to contain_exactly('status', 'start', 'stop', 'restart', 'help')
104
    end
105

106
    it 'filters subcommands by partial input' do
107
      # User typed: "mcp st<tab>" → words = ['mcp'], str = 'st'
108
      expect(plugin.cmd_mcp_tabs('st', ['mcp'])).to contain_exactly('status', 'start', 'stop')
109
    end
110

111
    it 'returns option completions for start subcommand' do
112
      # User typed: "mcp start <tab>" → words = ['mcp', 'start'], str = ''
113
      result = plugin.cmd_mcp_tabs('', ['mcp', 'start'])
114
      expect(result).to include('ServerHost=', 'RpcHost=', 'RpcPass=')
115
    end
116

117
    it 'returns empty array for status subcommand' do
118
      # User typed: "mcp status <tab>" → words = ['mcp', 'status'], str = ''
119
      expect(plugin.cmd_mcp_tabs('', ['mcp', 'status'])).to eq([])
120
    end
121
  end
122

123
  describe 'mcp status' do
124
    context 'when server is running with HTTP transport' do
125
      before do
126
        mcp_plugin.start_server({})
127
        reset_logging!
128
        allow(Time).to receive(:now).and_return(Time.at(1000))
129
        mcp_plugin.instance_variable_set(:@started_at, Time.at(1000))
130
      end
131

132
      it 'displays running state' do
133
        mcp_plugin.print_mcp_status
134
        expect(@output.join("\n")).to include('MCP server status: running')
135
      end
136

137
      it 'displays listening address and port' do
138
        mcp_plugin.print_mcp_status
139
        expect(@output.join("\n")).to include('http://localhost:3000')
140
      end
141

142
      it 'displays uptime' do
143
        mcp_plugin.print_mcp_status
144
        expect(@output.join("\n")).to include('Uptime:')
145
      end
146
    end
147

148
    context 'when server is stopped' do
149
      before do
150
        mcp_plugin.start_server({})
151
        mcp_plugin.stop_server
152
        reset_logging!
153
      end
154

155
      it 'displays stopped state' do
156
        mcp_plugin.print_mcp_status
157
        expect(@output.join("\n")).to include('MCP server status: stopped')
158
      end
159
    end
160

161
    context 'when server is running with custom host and port' do
162
      before do
163
        mcp_plugin.start_server('ServerHost' => '0.0.0.0', 'ServerPort' => '8080')
164
        reset_logging!
165
        allow(Time).to receive(:now).and_return(Time.at(1000))
166
        mcp_plugin.instance_variable_set(:@started_at, Time.at(1000))
167
      end
168

169
      it 'displays the custom listening address' do
170
        mcp_plugin.print_mcp_status
171
        expect(@output.join("\n")).to include('http://0.0.0.0:8080')
172
      end
173
    end
174

175
    context 'uptime formatting' do
176
      before do
177
        mcp_plugin.start_server({})
178
        reset_logging!
179
      end
180

181
      it 'formats seconds only' do
182
        allow(Time).to receive(:now).and_return(Time.at(1045))
183
        mcp_plugin.instance_variable_set(:@started_at, Time.at(1000))
184
        mcp_plugin.print_mcp_status
185
        expect(@output.join("\n")).to include('45s')
186
      end
187

188
      it 'formats minutes and seconds' do
189
        allow(Time).to receive(:now).and_return(Time.at(1125))
190
        mcp_plugin.instance_variable_set(:@started_at, Time.at(1000))
191
        mcp_plugin.print_mcp_status
192
        expect(@output.join("\n")).to include('2m 5s')
193
      end
194

195
      it 'formats hours, minutes, and seconds' do
196
        allow(Time).to receive(:now).and_return(Time.at(4661))
197
        mcp_plugin.instance_variable_set(:@started_at, Time.at(1000))
198
        mcp_plugin.print_mcp_status
199
        expect(@output.join("\n")).to include('1h 1m 1s')
200
      end
201
    end
202
  end
203

204
  describe 'mcp help' do
205
    it 'prints usage summary with all subcommands' do
206
      plugin.cmd_mcp_help
207
      combined = @output.join("\n")
208
      expect(combined).to include('status')
209
      expect(combined).to include('start')
210
      expect(combined).to include('stop')
211
      expect(combined).to include('restart')
212
      expect(combined).to include('help')
213
    end
214
  end
215

216
  describe '#cmd_mcp routing' do
217
    it 'routes to help when no subcommand given' do
218
      expect(plugin).to receive(:cmd_mcp_help)
219
      plugin.cmd_mcp
220
    end
221

222
    it 'routes to help for unrecognized subcommand' do
223
      expect(plugin).to receive(:cmd_mcp_help)
224
      plugin.cmd_mcp('unknown')
225
    end
226

227
    it 'routes to status subcommand' do
228
      expect(mcp_plugin).to receive(:print_mcp_status)
229
      plugin.cmd_mcp('status')
230
    end
231

232
    it 'routes to start subcommand' do
233
      expect(mcp_plugin).to receive(:start_server)
234
      plugin.cmd_mcp('start')
235
    end
236

237
    it 'routes to stop subcommand' do
238
      expect(mcp_plugin).to receive(:stop_server)
239
      plugin.cmd_mcp('stop')
240
    end
241

242
    it 'routes to restart subcommand' do
243
      expect(mcp_plugin).to receive(:restart_server)
244
      plugin.cmd_mcp('restart')
245
    end
246
  end
247

248
  describe 'mcp start' do
249
    context 'when server is stopped' do
250
      it 'starts the server successfully' do
251
        mcp_plugin.start_server({})
252
        expect(@output.join("\n")).to include('MCP server started')
253
      end
254

255
      it 'sets the server instance' do
256
        mcp_plugin.start_server({})
257
        expect(mcp_plugin.mcp_server).not_to be_nil
258
      end
259
    end
260

261
    context 'when server is already running' do
262
      before { mcp_plugin.start_server({}) }
263

264
      it 'prints an error that server is already running' do
265
        reset_logging!
266
        mcp_plugin.start_server({})
267
        expect(@error.join("\n")).to include('already running')
268
      end
269
    end
270
  end
271

272
  describe 'mcp stop' do
273
    context 'when server is running' do
274
      before do
275
        mcp_plugin.start_server({})
276
        reset_logging!
277
      end
278

279
      it 'stops the server and prints confirmation' do
280
        mcp_plugin.stop_server
281
        expect(@output.join("\n")).to include('MCP server stopped')
282
      end
283

284
      it 'clears the server instance' do
285
        mcp_plugin.stop_server
286
        expect(mcp_plugin.mcp_server).to be_nil
287
      end
288
    end
289

290
    context 'when server is already stopped' do
291
      it 'prints an error that server is already stopped' do
292
        mcp_plugin.stop_server
293
        expect(@error.join("\n")).to include('already stopped')
294
      end
295
    end
296
  end
297

298
  describe 'mcp restart' do
299
    context 'when server is running' do
300
      before { mcp_plugin.start_server({}) }
301

302
      it 'stops and then starts the server' do
303
        mcp_plugin.restart_server({})
304
        expect(mcp_plugin.mcp_server).not_to be_nil
305
      end
306

307
      it 'resets the started_at timestamp' do
308
        allow(Time).to receive(:now).and_return(Time.at(2000))
309
        mcp_plugin.restart_server({})
310
        expect(mcp_plugin.started_at).to eq(Time.at(2000))
311
      end
312
    end
313

314
    context 'when server is stopped' do
315
      it 'starts the server' do
316
        mcp_plugin.restart_server({})
317
        expect(mcp_plugin.mcp_server).not_to be_nil
318
      end
319
    end
320
  end
321
end
322

323