CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
# -*- coding: binary -*-
2
#
3
# Author:: Michael Neumann
4
# Copyright:: (c) 2005 by Michael Neumann
5
# License:: Same as Ruby's or BSD
6
#
7

8
require 'postgres_msf'
9
require 'postgres/buffer'
10
require 'rex/io/stream'
11

12
# Monkeypatch to preserve original code intent
13
# (postgres-pr originally defined read_exactly_n_bytes on IO
14
# as a while loop)
15
module Rex::IO::Stream
16
  def read_exactly_n_bytes(n)
17
    timed_read(n)
18
  end
19
end
20

21
# Namespace for Metasploit branch.
22
module Msf
23
module Db
24

25
module PostgresPR
26

27
class ParseError < RuntimeError; end
28
class DumpError < RuntimeError; end
29

30

31
# Base class representing a PostgreSQL protocol message
32
class Message
33
  # One character message-typecode to class map
34
  MsgTypeMap = Hash.new { UnknownMessageType }
35

36
  def self.register_message_type(type)
37
    raise "duplicate message type registration" if MsgTypeMap.has_key?(type)
38

39
    MsgTypeMap[type] = self
40

41
    self.const_set(:MsgType, type)
42
    class_eval "def message_type; MsgType end"
43
  end
44

45
  def self.read(stream, startup=false)
46
    type = stream.read_exactly_n_bytes(1) unless startup
47
    length = stream.read_exactly_n_bytes(4).to_s.unpack('N').first  # FIXME: length should be signed, not unsigned
48

49
    if length.nil?
50
      raise EOFError
51
    elsif length < 4
52
      raise ParseError
53
    end
54

55
    # If we didn't read any bytes and startup was not set, then type will be nil, so don't continue.
56
    unless startup
57
      if type.nil?
58
        return ParseError
59
      end
60
    end
61

62
    # initialize buffer
63
    buffer = Buffer.of_size(startup ? length : 1+length)
64
    buffer.write(type) unless startup
65
    buffer.write_int32_network(length)
66
    buffer.copy_from_stream(stream, length-4)
67

68
    (startup ? StartupMessage : MsgTypeMap[type]).create(buffer)
69
  end
70

71
  def self.create(buffer)
72
    obj = allocate
73
    obj.parse(buffer)
74
    obj
75
  end
76

77
  def self.dump(*args)
78
    new(*args).dump
79
  end
80

81
  def dump(body_size=0)
82
    buffer = Buffer.of_size(5 +  body_size)
83
    buffer.write(self.message_type)
84
    buffer.write_int32_network(4 + body_size)
85
    yield buffer if block_given?
86
    raise DumpError  unless buffer.at_end?
87
    return buffer.content
88
  end
89

90
  def parse(buffer)
91
    buffer.position = 5
92
    yield buffer if block_given?
93
    raise ParseError, buffer.inspect unless buffer.at_end?
94
  end
95

96
  def self.fields(*attribs)
97
    names = attribs.map {|name, type| name.to_s}
98
    arg_list = names.join(", ")
99
    ivar_list = names.map {|name| "@" + name }.join(", ")
100
    sym_list = names.map {|name| ":" + name }.join(", ")
101
    class_eval %[
102
      attr_accessor #{ sym_list }
103
      def initialize(#{ arg_list })
104
        #{ ivar_list } = #{ arg_list }
105
      end
106
    ]
107
  end
108
end
109

110
class UnknownMessageType < Message
111
  def dump
112
    raise
113
  end
114
end
115

116
class Authentication < Message
117
  register_message_type 'R'
118

119
  AuthTypeMap = {}
120

121
  def self.create(buffer)
122
    buffer.position = 5
123
    authtype = buffer.read_int32_network
124
    unless AuthTypeMap.key? authtype
125
      return UnknownAuthType.new(authtype, buffer)
126
    end
127

128
    klass = AuthTypeMap[authtype]
129
    obj = klass.allocate
130
    obj.parse(buffer)
131
    obj
132
  end
133

134
  def self.register_auth_type(type)
135
    raise "duplicate auth type registration" if AuthTypeMap.has_key?(type)
136
    AuthTypeMap[type] = self
137
    self.const_set(:AuthType, type)
138
    class_eval "def auth_type() AuthType end"
139
  end
140

141
  # the dump method of class Message
142
  alias message__dump dump
143

144
  def dump
145
    super(4) do |buffer|
146
      buffer.write_int32_network(self.auth_type)
147
    end
148
  end
149

150
  def parse(buffer)
151
    super do
152
      auth_t = buffer.read_int32_network
153
      raise ParseError unless auth_t == self.auth_type
154
      yield if block_given?
155
    end
156
  end
157
end
158

159
class UnknownAuthType < Authentication
160
  attr_reader :auth_type
161
  attr_reader :buffer
162

163
  def initialize(auth_type, buffer)
164
    @auth_type = auth_type
165
    @buffer = buffer
166
  end
167
end
168

169
class AuthentificationOk < Authentication
170
  register_auth_type 0
171
end
172

173
class AuthentificationKerberosV4 < Authentication
174
  register_auth_type 1
175
end
176

177
class AuthentificationKerberosV5 < Authentication
178
  register_auth_type 2
179
end
180

181
class AuthentificationClearTextPassword < Authentication
182
  register_auth_type 3
183
end
184

185
module SaltedAuthentificationMixin
186
  attr_accessor :salt
187

188
  def initialize(salt)
189
    @salt = salt
190
  end
191

192
  def dump
193
    raise DumpError unless @salt.size == self.salt_size
194

195
    message__dump(4 + self.salt_size) do |buffer|
196
      buffer.write_int32_network(self.auth_type)
197
      buffer.write(@salt)
198
    end
199
  end
200

201
  def parse(buffer)
202
    super do
203
      @salt = buffer.read(self.salt_size)
204
    end
205
  end
206
end
207

208
class AuthentificationCryptPassword < Authentication
209
  register_auth_type 4
210
  include SaltedAuthentificationMixin
211
  def salt_size; 2 end
212
end
213

214

215
class AuthentificationMD5Password < Authentication
216
  register_auth_type 5
217
  include SaltedAuthentificationMixin
218
  def salt_size; 4 end
219
end
220

221
class AuthentificationSCMCredential < Authentication
222
  register_auth_type 6
223
end
224

225
# SASL Overview
226
# https://www.postgresql.org/docs/current/sasl-authentication.html
227
#
228
# Binary format:
229
# https://www.postgresql.org/docs/current/protocol-message-formats.html
230
class AuthenticationSASL < Authentication
231
  # Int32(10) - Specifies that SASL authentication is required.
232
  register_auth_type 10
233

234
  # @return [Array<String>] Name of a SASL authentication mechanisms
235
  attr_reader :mechanisms
236

237
  # @param [Array<String>] mechanisms
238
  def initialize(mechanisms: [])
239
    @mechanisms = mechanisms
240
  end
241

242
  def dump
243
    auth_type_byte_size = 4
244
    mechanism_bytes_size = mechanisms.sum(&:size) + (mechanisms.size + 1)
245
    message__dump(auth_type_byte_size + mechanism_bytes_size) do |buffer|
246
      buffer.write_int32_network(self.auth_type)
247
      mechanisms.each do |mechanism|
248
        buffer.write_cstring(mechanism)
249
      end
250
      buffer.write(Buffer::NUL)
251
    end
252
  end
253

254
  def parse(buffer)
255
    super do
256
      # The message body is a list of SASL authentication mechanisms, in the
257
      # server's order of preference. A zero byte is required as terminator after
258
      # the last authentication mechanism name.
259
      # https://github.com/postgres/postgres/blob/74a2dfee2255a1bace9b0053d014c4efa2823f4d/doc/src/sgml/protocol.sgml#L3584-L3602
260
      @mechanisms ||= []
261
      while buffer.peek != Buffer::NUL
262
        @mechanisms << buffer.read_cstring
263
      end
264
      _null = buffer.read(1)
265
    end
266
  end
267
end
268

269
# AuthenticationSASLContinue (B)
270
# https://www.postgresql.org/docs/current/protocol-message-formats.html
271
class AuthenticationSASLContinue < Authentication
272
  # Int32(11) - Specifies that this message contains a SASL challenge.
273
  register_auth_type 11
274

275
  # @return [String] SASL data, specific to the SASL mechanism being used.
276
  attr_reader :value
277

278
  # @param [String, nil] value
279
  def initialize(value: nil)
280
    @value = value
281
  end
282

283
  def dump
284
    auth_type_byte_size = 4
285
    value_size = value.size
286
    message__dump(auth_type_byte_size + value_size) do |buffer|
287
      buffer.write_int32_network(self.auth_type)
288
      buffer.write(value)
289
    end
290
  end
291

292
  def parse(buffer)
293
    super do
294
      @value = buffer.read_rest
295
    end
296
  end
297
end
298

299
# AuthenticationSASLFinal (B)
300
# https://www.postgresql.org/docs/current/protocol-message-formats.html
301
class AuthenticationSASLFinal < Authentication
302
  # Int32(11) - Specifies that this message contains a SASL challenge.
303
  register_auth_type 12
304

305
  # @return [String] SASL outcome "additional data", specific to the SASL mechanism being used.
306
  attr_reader :value
307

308
  # @param [String] value
309
  def initialize(value:)
310
    @value = value
311
  end
312

313
  def dump
314
    auth_type_byte_size = 4
315
    value_size = value.size
316
    message__dump(auth_type_byte_size + value_size) do |buffer|
317
      buffer.write_int32_network(self.auth_type)
318
      buffer.write(value)
319
    end
320
  end
321

322
  def parse(buffer)
323
    super do
324
      @value = buffer.read_rest
325
    end
326
  end
327
end
328

329
class PasswordResponseMessage < Message
330
  # Identifies the message as a password response. Note that this is also used for GSSAPI, SSPI and SASL response messages.
331
  # The exact message type can be deduced from the context.
332
  register_message_type 'p'
333
end
334

335
class PasswordMessage < PasswordResponseMessage
336
  fields :password
337

338
  def dump
339
    super(@password.size + 1) do |buffer|
340
      buffer.write_cstring(@password)
341
    end
342
  end
343

344
  def parse(buffer)
345
    super do
346
      @password = buffer.read_cstring
347
    end
348
  end
349
end
350

351
# SASLInitialResponse (F). The client sends a SASLInitialResponse after choosing a SASL mechanism. The message includes the name of the selected
352
# mechanism, and an optional Initial Client Response, if the selected mechanism uses that.
353
#
354
# https://www.postgresql.org/docs/current/protocol-message-formats.html
355
# https://www.postgresql.org/docs/current/sasl-authentication.html
356
class SaslInitialResponseMessage < PasswordResponseMessage
357

358
  # @return [String] Name of the SASL authentication mechanism that the client selected.
359
  attr_reader :mechanism
360

361
  # @return [String] SASL mechanism specific "Initial Response" - specific to the SASL mechanism used
362
  attr_reader :value
363

364
  # @param [String] mechanism
365
  # @param [String] value
366
  def initialize(mechanism: nil, value: nil)
367
    @mechanism = mechanism
368
    @value = value
369
  end
370

371
  def dump
372
    mechanism_size = mechanism.size + Buffer::NUL.size
373
    value_size_prefix_size = 4
374
    value_size = value.size
375
    message_size = mechanism_size + value_size_prefix_size + value_size
376

377
    super(message_size) do |buffer|
378
      buffer.write_cstring(mechanism)
379
      buffer.write_int32_network(value_size)
380
      buffer.write(value)
381
    end
382
  end
383

384
  def parse(buffer)
385
    super do
386
      @mechanism = buffer.read_cstring
387
      _value_size_prefix_size = buffer.read_int32_network
388
      @value = buffer.read_rest
389
    end
390
  end
391

392
  def ==(other)
393
    self.class == other.class &&
394
      @mechanism == other.mechanism &&
395
      @value == other.value
396
  end
397
end
398

399
# SASLResponse (F)
400
#
401
# https://www.postgresql.org/docs/current/protocol-message-formats.html
402
# https://www.postgresql.org/docs/current/sasl-authentication.html
403
class SASLResponseMessage < PasswordResponseMessage
404

405
  # @return [String] SASL mechanism specific "Initial Response" - specific to the SASL mechanism used
406
  attr_reader :value
407

408
  # @param [String] value
409
  def initialize(value: nil)
410
    @value = value
411
  end
412

413
  def dump
414
    super(value.size) do |buffer|
415
      buffer.write(value)
416
    end
417
  end
418

419
  def parse(buffer)
420
    super do
421
      @value = buffer.read_rest
422
    end
423
  end
424

425
  def ==(other)
426
    self.class == other.class &&
427
      @value == other.value
428
  end
429
end
430

431
class ParameterStatus < Message
432
  register_message_type 'S'
433
  fields :key, :value
434

435
  def dump
436
    super(@key.size + 1 + @value.size + 1) do |buffer|
437
      buffer.write_cstring(@key)
438
      buffer.write_cstring(@value)
439
    end
440
  end
441

442
  def parse(buffer)
443
    super do
444
      @key = buffer.read_cstring
445
      @value = buffer.read_cstring
446
    end
447
  end
448
end
449

450
class BackendKeyData < Message
451
  register_message_type 'K'
452
  fields :process_id, :secret_key
453

454
  def dump
455
    super(4 + 4) do |buffer|
456
      buffer.write_int32_network(@process_id)
457
      buffer.write_int32_network(@secret_key)
458
    end
459
  end
460

461
  def parse(buffer)
462
    super do
463
      @process_id = buffer.read_int32_network
464
      @secret_key = buffer.read_int32_network
465
    end
466
  end
467
end
468

469
class ReadyForQuery < Message
470
  register_message_type 'Z'
471
  fields :backend_transaction_status_indicator
472

473
  def dump
474
    super(1) do |buffer|
475
      buffer.write_byte(@backend_transaction_status_indicator)
476
    end
477
  end
478

479
  def parse(buffer)
480
    super do
481
      @backend_transaction_status_indicator = buffer.read_byte
482
    end
483
  end
484
end
485

486
class DataRow < Message
487
  register_message_type 'D'
488
  fields :columns
489

490
  def dump
491
    sz = @columns.inject(2) {|sum, col| sum + 4 + (col ? col.size : 0)}
492
    super(sz) do |buffer|
493
      buffer.write_int16_network(@columns.size)
494
      @columns.each {|col|
495
        buffer.write_int32_network(col ? col.size : -1)
496
        buffer.write(col) if col
497
      }
498
    end
499
  end
500

501
  def parse(buffer)
502
    super do
503
      n_cols = buffer.read_int16_network
504
      @columns = (1..n_cols).collect {
505
        len = buffer.read_int32_network
506
        if len == -1
507
          nil
508
        else
509
          buffer.read(len)
510
        end
511
      }
512
    end
513
  end
514
end
515

516
class CommandComplete < Message
517
  register_message_type 'C'
518
  fields :cmd_tag
519

520
  def dump
521
    super(@cmd_tag.size + 1) do |buffer|
522
      buffer.write_cstring(@cmd_tag)
523
    end
524
  end
525

526
  def parse(buffer)
527
    super do
528
      @cmd_tag = buffer.read_cstring
529
    end
530
  end
531
end
532

533
class EmptyQueryResponse < Message
534
  register_message_type 'I'
535
end
536

537
module NoticeErrorMixin
538
  attr_accessor :field_type, :field_values
539

540
  def initialize(field_type=0, field_values=[])
541
    raise ArgumentError if field_type == 0 and not field_values.empty?
542
    @field_type, @field_values = field_type, field_values
543
  end
544

545
  def dump
546
    raise ArgumentError if @field_type == 0 and not @field_values.empty?
547

548
    sz = 1
549
    sz += @field_values.inject(1) {|sum, fld| sum + fld.size + 1} unless @field_type == 0
550

551
    super(sz) do |buffer|
552
      buffer.write_byte(@field_type)
553
      break if @field_type == 0
554
      @field_values.each {|fld| buffer.write_cstring(fld) }
555
      buffer.write_byte(0)
556
    end
557
  end
558

559
  def parse(buffer)
560
    super do
561
      @field_type = buffer.read_byte
562
      break if @field_type == 0
563
      @field_values = []
564
      while buffer.position < buffer.size-1
565
        @field_values << buffer.read_cstring
566
      end
567
      terminator = buffer.read_byte
568
      raise ParseError unless terminator == 0
569
    end
570
  end
571
end
572

573
class NoticeResponse < Message
574
  register_message_type 'N'
575
  include NoticeErrorMixin
576
end
577

578
class ErrorResponse < Message
579
  register_message_type 'E'
580
  include NoticeErrorMixin
581
end
582

583
# TODO
584
class CopyInResponse < Message
585
  register_message_type 'G'
586
end
587

588
# TODO
589
class CopyOutResponse < Message
590
  register_message_type 'H'
591
end
592

593
class Parse < Message
594
  register_message_type 'P'
595
  fields :query, :stmt_name, :parameter_oids
596

597
  def initialize(query, stmt_name="", parameter_oids=[])
598
    @query, @stmt_name, @parameter_oids = query, stmt_name, parameter_oids
599
  end
600

601
  def dump
602
    sz = @stmt_name.size + 1 + @query.size + 1 + 2 + (4 * @parameter_oids.size)
603
    super(sz) do |buffer|
604
      buffer.write_cstring(@stmt_name)
605
      buffer.write_cstring(@query)
606
      buffer.write_int16_network(@parameter_oids.size)
607
      @parameter_oids.each {|oid| buffer.write_int32_network(oid) }
608
    end
609
  end
610

611
  def parse(buffer)
612
    super do
613
      @stmt_name = buffer.read_cstring
614
      @query = buffer.read_cstring
615
      n_oids = buffer.read_int16_network
616
      @parameter_oids = (1..n_oids).collect {
617
        # TODO: zero means unspecified. map to nil?
618
        buffer.read_int32_network
619
      }
620
    end
621
  end
622
end
623

624
class ParseComplete < Message
625
  register_message_type '1'
626
end
627

628
class Query < Message
629
  register_message_type 'Q'
630
  fields :query
631

632
  def dump
633
    super(@query.size + 1) do |buffer|
634
      buffer.write_cstring(@query)
635
    end
636
  end
637

638
  def parse(buffer)
639
    super do
640
      @query = buffer.read_cstring
641
    end
642
  end
643
end
644

645
class RowDescription < Message
646
  register_message_type 'T'
647
  fields :fields
648

649
  class FieldInfo < Struct.new(:name, :oid, :attr_nr, :type_oid, :typlen, :atttypmod, :formatcode); end
650

651
  def dump
652
    sz = @fields.inject(2) {|sum, fld| sum + 18 + fld.name.size + 1 }
653
    super(sz) do |buffer|
654
      buffer.write_int16_network(@fields.size)
655
      @fields.each { |f|
656
        buffer.write_cstring(f.name)
657
        buffer.write_int32_network(f.oid)
658
        buffer.write_int16_network(f.attr_nr)
659
        buffer.write_int32_network(f.type_oid)
660
        buffer.write_int16_network(f.typlen)
661
        buffer.write_int32_network(f.atttypmod)
662
        buffer.write_int16_network(f.formatcode)
663
      }
664
    end
665
  end
666

667
  def parse(buffer)
668
    super do
669
      n_fields = buffer.read_int16_network
670
      @fields = (1..n_fields).collect {
671
        f = FieldInfo.new
672
        f.name       = buffer.read_cstring
673
        f.oid        = buffer.read_int32_network
674
        f.attr_nr    = buffer.read_int16_network
675
        f.type_oid   = buffer.read_int32_network
676
        f.typlen     = buffer.read_int16_network
677
        f.atttypmod  = buffer.read_int32_network
678
        f.formatcode = buffer.read_int16_network
679
        f
680
      }
681
    end
682
  end
683
end
684

685
class StartupMessage < Message
686
  fields :proto_version, :params
687

688
  def dump
689
    sz = @params.inject(4 + 4) {|sum, kv| sum + kv[0].size + 1 + kv[1].size + 1} + 1
690

691
    buffer = Buffer.of_size(sz)
692
    buffer.write_int32_network(sz)
693
    buffer.write_int32_network(@proto_version)
694
    @params.each_pair {|key, value|
695
      buffer.write_cstring(key)
696
      buffer.write_cstring(value)
697
    }
698
    buffer.write_byte(0)
699

700
    raise DumpError unless buffer.at_end?
701
    return buffer.content
702
  end
703

704
  def parse(buffer)
705
    buffer.position = 4
706

707
    @proto_version = buffer.read_int32_network
708
    @params = {}
709

710
    while buffer.position < buffer.size-1
711
      key = buffer.read_cstring
712
      val = buffer.read_cstring
713
      @params[key] = val
714
    end
715

716
    nul = buffer.read_byte
717
    raise ParseError unless nul == 0
718
    raise ParseError unless buffer.at_end?
719
  end
720
end
721

722
class SSLRequest < Message
723
  fields :ssl_request_code
724

725
  def dump
726
    sz = 4 + 4
727
    buffer = Buffer.of_size(sz)
728
    buffer.write_int32_network(sz)
729
    buffer.write_int32_network(@ssl_request_code)
730
    raise DumpError unless buffer.at_end?
731
    return buffer.content
732
  end
733

734
  def parse(buffer)
735
    buffer.position = 4
736
    @ssl_request_code = buffer.read_int32_network
737
    raise ParseError unless buffer.at_end?
738
  end
739
end
740

741
=begin
742
# TODO: duplicate message-type, split into client/server messages
743
class Sync < Message
744
  register_message_type 'S'
745
end
746
=end
747

748
class Terminate < Message
749
  register_message_type 'X'
750
end
751

752
end # module PostgresPR
753

754
end
755
end
756

757