Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
# -*- coding: binary -*-
2

3
require 'bindata'
4

5
module Rex::Proto::Http::WebSocket::AmazonSsm
6
  module PayloadType
7
    Output               = 1
8
    Error                = 2
9
    Size                 = 3
10
    Parameter            = 4
11
    HandshakeRequest     = 5
12
    HandshakeResponse    = 6
13
    HandshakeComplete    = 7
14
    EncChallengeRequest  = 8
15
    EncChallengeResponse = 9
16
    Flag                 = 10
17

18
    def self.from_val(v)
19
      self.constants.find {|c| self.const_get(c) == v }
20
    end
21
  end
22

23
  module UUID
24
    def self.unpack(bbuf)
25
      sbuf = ""
26
      [8...12].each do |idx|
27
        sbuf << Rex::Text.to_hex(bbuf[idx])
28
      end
29
      sbuf << '-'
30
      [12...14].each do |idx|
31
        sbuf << Rex::Text.to_hex(bbuf[idx])
32
      end
33
      sbuf << '-'
34
      [14...16].each do |idx|
35
        sbuf << Rex::Text.to_hex(bbuf[idx])
36
      end
37
      sbuf << '-'
38
      [0...2].each do |idx|
39
        sbuf << Rex::Text.to_hex(bbuf[idx])
40
      end
41
      sbuf << '-'
42
      [2...8].each do |idx|
43
        sbuf << Rex::Text.to_hex(bbuf[idx])
44
      end
45
      sbuf.gsub("\\x",'')
46
    end
47

48
    def self.pack(sbuf)
49
      parts = sbuf.split('-').map do |seg|
50
        seg.chars.each_slice(2).map {|e| "\\x#{e.join}"}.join
51
      end
52
      [3, 4, 0, 1, 2].map do |part|
53
        Rex::Text.hex_to_raw(parts[part])
54
      end.join
55
    end
56

57
    def self.rand
58
      self.unpack(Rex::Text.rand_text(16))
59
    end
60
  end
61

62
  module Interface
63
    module SsmChannelMethods
64
      attr_accessor :rows
65
      attr_accessor :cols
66

67
      def _start_ssm_keepalive
68
        @keepalive_thread = Rex::ThreadFactory.spawn('SsmChannel-Keepalive', false) do
69
          while not closed? or @websocket.closed?
70
            write ''
71
            Rex::ThreadSafe.sleep(::Random.rand * 10 + 15)
72
          end
73
          @keepalive_thread = nil
74
        end
75
      end
76

77
      def close
78
        @keepalive_thread.kill if @keepalive_thread
79
        @keepalive_thread = nil
80
        super
81
      end
82

83
      def acknowledge_output(output_frame)
84
        ack = output_frame.to_ack
85
        # ack.header.sequence_number = @out_seq_num
86
        @websocket.put_wsbinary(ack.to_binary_s)
87
        # wlog("SsmChannel: acknowledge output #{output_frame.uuid}")
88
        output_frame.uuid
89
      end
90

91
      def pause_publication
92
        msg = SsmFrame.create_pause_pub
93
        @publication = false
94
        @websocket.put_wsbinary(msg.to_binary_s)
95
      end
96

97
      def start_publication
98
        msg = SsmFrame.create_start_pub
99
        @publication = true
100
        @websocket.put_wsbinary(msg.to_binary_s)
101
      end
102

103
      def handle_output_data(output_frame)
104
        return nil if @ack_message == output_frame.uuid
105

106
        @ack_message = acknowledge_output(output_frame)
107
        # TODO: handle Payload::* types
108
        if ![PayloadType::Output, PayloadType::Error].any? { |e| e == output_frame.payload_type }
109
          wlog("SsmChannel got unhandled output payload type: #{Payload.from_val(output_frame.payload_type)}")
110
          return nil
111
        end
112

113
        output_frame.payload_data.value
114
      end
115

116
      def handle_acknowledge(ack_frame)
117
        # wlog("SsmChannel: got acknowledge message #{ack_frame.uuid}")
118
        begin
119
          seq_num = JSON.parse(ack_frame.payload_data)['AcknowledgedMessageSequenceNumber'].to_i
120
          @ack_seq_num = seq_num if seq_num > @ack_seq_num
121
        rescue => e
122
          elog("SsmChannel failed to parse ack JSON #{ack_frame.payload_data} due to #{e}!")
123
        end
124
        nil
125
      end
126

127
      def update_term_size
128
        return unless ::IO.console
129

130
        rows, cols = ::IO.console.winsize
131
        unless rows == self.rows && cols == self.cols
132
          set_term_size(cols, rows)
133
          self.rows = rows
134
          self.cols = cols
135
        end
136
      end
137

138
      def set_term_size(cols, rows)
139
        data = JSON.generate({cols: cols, rows: rows})
140
        frame = SsmFrame.create(data)
141
        frame.payload_type = PayloadType::Size
142
        @websocket.put_wsbinary(frame.to_binary_s)
143
      end
144
    end
145

146
    class SsmChannel < Rex::Proto::Http::WebSocket::Interface::Channel
147
      include SsmChannelMethods
148
      attr_reader :run_ssm_pub, :out_seq_num, :ack_seq_num, :ack_message
149

150
      def initialize(websocket)
151
        @ack_seq_num = 0
152
        @out_seq_num = 0
153
        @run_ssm_pub = true
154
        @ack_message = nil
155
        @publication = false
156

157
        super(websocket, write_type: :binary)
158
      end
159

160
      def on_data_read(data, _data_type)
161
        return data if data.blank?
162

163
        ssm_frame = SsmFrame.read(data)
164
        case ssm_frame.header.message_type.strip
165
        when 'output_stream_data'
166
          @publication = true # Linux sends stream data before sending start_publication message
167
          return handle_output_data(ssm_frame)
168
        when 'acknowledge'
169
          # update ACK seqno
170
          handle_acknowledge(ssm_frame)
171
        when 'start_publication'
172
          @out_seq_num = @ack_seq_num if @out_seq_num > 0
173
          @publication = true
174
          # handle session resumption - foregrounding or resumption of input
175
        when 'pause_publication'
176
          # @websocket.put_wsbinary(ssm_frame.to_ack.to_binary_s)
177
          @publication = false
178
          # handle session suspension - backgrounding or general idle
179
        when 'input_stream_data'
180
          # this is supposed to be a one way street
181
          emsg = "SsmChannel received input_stream_data from SSM (!!)"
182
          elog(emsg)
183
          raise emsg
184
        when 'channel_closed'
185
          elog("SsmChannel got closed message #{ssm_frame.uuid}")
186
          close
187
        else
188
          raise Rex::Proto::Http::WebSocket::ConnectionError.new(
189
            msg: "Unknown AWS SSM message type: #{ssm_frame.header.message_type}"
190
          )
191
        end
192

193
        nil
194
      end
195

196
      def on_data_write(data)
197
        start_publication if not @publication
198
        frame = SsmFrame.create(data)
199
        frame.header.sequence_number = @out_seq_num
200
        @out_seq_num += 1
201
        frame.to_binary_s
202
      end
203

204
      def publishing?
205
        @publication
206
      end
207
    end
208

209
    def to_ssm_channel(publish_timeout: 10)
210
      chan = SsmChannel.new(self)
211

212
      if publish_timeout
213
        # Waiting for the channel to start publishing
214
        (publish_timeout * 2).times do
215
          break if chan.publishing?
216

217
          sleep 0.5
218
        end
219

220
        raise Rex::TimeoutError.new('Timed out while waiting for the channel to start publishing.') unless chan.publishing?
221
      end
222

223
      chan
224
    end
225
  end
226

227
  class SsmFrame < BinData::Record
228
    endian :big
229

230
    struct :header do
231
      endian :big
232

233
      uint32 :header_length, initial_value: 116
234
      string :message_type, length: 32, pad_byte: 0x20, initial_value: 'input_stream_data'
235
      uint32 :schema_version, initial_value: 1
236
      uint64 :created_date, default_value: lambda { (Time.now.to_f * 1000).to_i }
237
      uint64 :sequence_number, initial_value: 0
238
      uint64 :flags, value: 0 #lambda { sequence_number == 0 ? 1 : 0 }
239
      string :message_id, length: 16, initial_value: UUID.pack(UUID.rand)
240
    end
241

242
    string :payload_digest, length: 32, default_value: -> { Digest::SHA256.digest(payload_data) }
243
    uint32 :payload_type, default_value: PayloadType::Output
244
    uint32 :payload_length, value: -> { payload_data.length }
245
    string :payload_data, read_length: -> { payload_length }
246

247
    class << self
248
      def create(data = nil, mtype = 'input_stream_data')
249
        return data if data.is_a?(SsmFrame)
250

251
        frame = SsmFrame.new(header: {
252
          message_type: mtype,
253
          created_date: (Time.now.to_f * 1000).to_i,
254
          message_id: UUID.pack(UUID.rand)
255
        })
256
        if !data.nil?
257
          frame.payload_data   = data
258
          frame.payload_digest = Digest::SHA256.digest(data)
259
          frame.payload_length = data.length
260
          frame.payload_type   = PayloadType::Output
261
        end
262
        frame
263
      end
264

265
      def create_pause_pub
266
        uuid = UUID.rand
267
        time = Time.now
268
        data = JSON.generate({
269
          MessageType: 'pause_publication',
270
          SchemaVersion: 1,
271
          MessageId: uuid,
272
          CreateData: time.strftime("%Y-%m-%dT%T.%LZ")
273
        })
274
        frame = SsmFrame.new( header: {
275
          message_type: 'pause_publication',
276
          created_date: (time.to_f * 1000).to_i,
277
          message_id: UUID.pack(uuid)
278
        })
279
        frame.payload_data   = data
280
        frame.payload_digest = Digest::SHA256.digest(data)
281
        frame.payload_length = data.length
282
        frame.payload_type   = 0
283
        frame
284
      end
285

286
      def create_start_pub
287
        data = 'start_publication'
288
        frame = SsmFrame.new( header: {
289
          message_type: data,
290
          created_date: (Time.now.to_f * 1000).to_i,
291
          message_id: UUID.pack(UUID.rand)
292
        })
293
        frame.payload_data   = data
294
        frame.payload_digest = Digest::SHA256.digest(data)
295
        frame.payload_length = data.length
296
        frame.payload_type   = 0
297
        frame
298
      end
299

300
      def from_ws_frame(wsframe)
301
        SsmFrame.read(wsframe.payload_data)
302
      end
303
    end
304

305
    def uuid
306
      UUID.unpack(header.message_id)
307
    end
308

309
    def to_ack
310
      data = JSON.generate({
311
        AcknowledgedMessageType: header.message_type.strip,
312
        AcknowledgedMessageId: uuid,
313
        AcknowledgedMessageSequenceNumber: header.sequence_number.to_i,
314
        IsSequentialMessage: true
315
      })
316
      ack = SsmFrame.create(data, 'acknowledge')
317
      ack.header.sequence_number = header.sequence_number
318
      ack.header.flags = header.flags
319
      ack
320
    end
321

322
    def length
323
      to_binary_s.length
324
    end
325
  end
326
  #
327
  # Initiates a WebSocket session based on the params of SSM::Client#start_session
328
  #
329
  # @param [Aws::SSM::Types::StartSessionResponse] :session_init Parameters returned by #start_session
330
  # @param [Integer] :timeout
331
  #
332
  # @return [Socket] Socket representing the authenticates SSM WebSocket connection
333
  def connect_ssm_ws(session_init, timeout = 20)
334
    # hack-up a "graceful fail-down" in the caller
335
    # raise Rex::Proto::Http::WebSocket::ConnectionError.new(msg: 'WebSocket sessions still need structs/parsing')
336
    ws_key = session_init.token_value
337
    ssm_id = session_init.session_id
338
    ws_url = URI.parse(session_init.stream_url)
339
    opts   = {}
340
    opts['vhost']   = ws_url.host
341
    opts['uri']     = ws_url.to_s.sub(/^.*#{ws_url.host}/, '')
342
    opts['headers'] = {
343
      'Connection'            => 'Upgrade',
344
      'Upgrade'               => 'WebSocket',
345
      'Sec-WebSocket-Version' => 13,
346
      'Sec-WebSocket-Key'     => ws_key
347
    }
348
    ctx = {
349
      'Msf'        => framework,
350
      'MsfExploit' => self
351
    }
352
    http_client = Rex::Proto::Http::Client.new(ws_url.host, 443, ctx, true)
353
    raise Rex::Proto::Http::WebSocket::ConnectionError.new if http_client.nil?
354

355
    # Send upgrade request
356
    req = http_client.request_raw(opts)
357
    res = http_client.send_recv(req, timeout)
358
    # Verify upgrade
359
    unless res&.code == 101
360
      http_client.close
361
      raise Rex::Proto::Http::WebSocket::ConnectionError.new(http_response: res)
362
    end
363
    # see: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Sec-WebSocket-Accept
364
    accept_ws_key = Rex::Text.encode_base64(OpenSSL::Digest::SHA1.digest(ws_key + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'))
365
    unless res.headers['Sec-WebSocket-Accept'] == accept_ws_key
366
      http_client.close
367
      raise Rex::Proto::Http::WebSocket::ConnectionError.new(msg: 'Invalid Sec-WebSocket-Accept header', http_response: res)
368
    end
369
    # Extract and extend connection object
370
    socket = http_client.conn
371
    socket.extend(Rex::Proto::Http::WebSocket::Interface)
372
    # Send initialization handshake
373
    ssm_wsock_init = JSON.generate({
374
      MessageSchemaVersion: '1.0',
375
      RequestId: UUID.rand,
376
      TokenValue: ws_key
377
    })
378
    socket.put_wstext(ssm_wsock_init)
379
    # Extend with interface
380
    socket.extend(Interface)
381
  end
382
end
383

384