1
# frozen_string_literal: true
2

3
require 'msf/core/mcp'
4
require 'webmock/rspec'
5

6
RSpec.describe 'MessagePack Re-Authentication Flow Integration' do
7
  before(:all) do
8
    WebMock.disable_net_connect!(allow_localhost: false)
9
  end
10

11
  after(:all) do
12
    WebMock.allow_net_connect!
13
  end
14

15
  let(:host) { 'localhost' }
16
  let(:port) { 55553 }
17
  let(:endpoint) { '/api/' }
18
  let(:api_url) { "https://#{host}:#{port}#{endpoint}" }
19
  let(:user) { 'test_user' }
20
  let(:password) { 'test_password' }
21

22
  describe 'Automatic Re-Authentication on Token Expiry' do
23
    it 're-authenticates and retries when API call returns 401' do
24
      call_count = 0
25

26
      # Stub all POST requests and dispatch based on body content
27
      stub_request(:post, api_url).to_return do |request|
28
        body = MessagePack.unpack(request.body)
29
        call_count += 1
30

31
        case call_count
32
        when 1
33
          # Initial authentication succeeds
34
          expect(body[0]).to eq('auth.login')
35
          {
36
            status: 200,
37
            body: { 'result' => 'success', 'token' => 'initial_token' }.to_msgpack,
38
            headers: { 'Content-Type' => 'binary/message-pack' }
39
          }
40
        when 2
41
          # First API call returns 401 (token expired)
42
          expect(body[0]).to eq('module.search')
43
          expect(body[1]).to eq('initial_token')
44
          {
45
            status: 401,
46
            body: { 'error_message' => 'Token expired' }.to_msgpack,
47
            headers: { 'Content-Type' => 'binary/message-pack' }
48
          }
49
        when 3
50
          # Re-authentication succeeds with new token
51
          expect(body[0]).to eq('auth.login')
52
          {
53
            status: 200,
54
            body: { 'result' => 'success', 'token' => 'refreshed_token' }.to_msgpack,
55
            headers: { 'Content-Type' => 'binary/message-pack' }
56
          }
57
        when 4
58
          # Retry with new token succeeds
59
          expect(body[0]).to eq('module.search')
60
          expect(body[1]).to eq('refreshed_token')
61
          {
62
            status: 200,
63
            body: [{ 'fullname' => 'exploit/test', 'type' => 'exploit', 'name' => 'test' }].to_msgpack,
64
            headers: { 'Content-Type' => 'binary/message-pack' }
65
          }
66
        else
67
          raise "Unexpected request ##{call_count}: #{body.inspect}"
68
        end
69
      end
70

71
      client = Msf::MCP::Metasploit::MessagePackClient.new(
72
        host: host,
73
        port: port,
74
        endpoint: endpoint
75
      )
76
      client.authenticate(user, password)
77

78
      # This call should trigger: 401 → re-auth → retry → success
79
      result = client.search_modules('smb')
80

81
      expect(result).to be_an(Array)
82
      expect(result.first['fullname']).to eq('exploit/test')
83
      expect(client.instance_variable_get(:@token)).to eq('refreshed_token')
84
      expect(call_count).to eq(4)
85
    end
86

87
    it 'gives up after max retries when re-auth succeeds but API keeps failing' do
88
      call_count = 0
89

90
      stub_request(:post, api_url).to_return do |request|
91
        body = MessagePack.unpack(request.body)
92
        call_count += 1
93

94
        if body[0] == 'auth.login'
95
          {
96
            status: 200,
97
            body: { 'result' => 'success', 'token' => "token_#{call_count}" }.to_msgpack,
98
            headers: { 'Content-Type' => 'binary/message-pack' }
99
          }
100
        else
101
          # API calls always return 401
102
          {
103
            status: 401,
104
            body: { 'error_message' => 'Token invalid' }.to_msgpack,
105
            headers: { 'Content-Type' => 'binary/message-pack' }
106
          }
107
        end
108
      end
109

110
      client = Msf::MCP::Metasploit::MessagePackClient.new(
111
        host: host,
112
        port: port,
113
        endpoint: endpoint
114
      )
115
      client.authenticate(user, password)
116

117
      # Should exhaust retries (max_retries=2) and re-raise
118
      expect {
119
        client.search_modules('smb')
120
      }.to raise_error(Msf::MCP::Metasploit::AuthenticationError)
121
    end
122

123
    it 'propagates re-auth failure through the tool layer as an error response' do
124
      stub_request(:post, api_url)
125
        .with(body: ['auth.login', user, password].to_msgpack)
126
        .to_return(
127
          status: 200,
128
          body: { 'result' => 'success', 'token' => 'test_token' }.to_msgpack,
129
          headers: { 'Content-Type' => 'binary/message-pack' }
130
        )
131

132
      # All subsequent requests return 401
133
      stub_request(:post, api_url)
134
        .with { |req| MessagePack.unpack(req.body)[0] != 'auth.login' }
135
        .to_return(
136
          status: 401,
137
          body: { 'error_message' => 'Token invalid' }.to_msgpack,
138
          headers: { 'Content-Type' => 'binary/message-pack' }
139
        )
140

141
      client = Msf::MCP::Metasploit::MessagePackClient.new(
142
        host: host,
143
        port: port,
144
        endpoint: endpoint
145
      )
146
      client.authenticate(user, password)
147

148
      limiter = Msf::MCP::Security::RateLimiter.new(requests_per_minute: 60)
149
      server_context = { msf_client: client, rate_limiter: limiter }
150

151
      result = Msf::MCP::Tools::SearchModules.call(query: 'smb', server_context: server_context)
152

153
      expect(result.error?).to be true
154
      expect(result.content.first[:text]).to include('Authentication failed')
155
    end
156
  end
157
end
158

159