metasploit-gs/lib/rbmysql/protocol.rb

# coding: ascii-8bit
# Copyright (C) 2008-2012 TOMITA Masahiro
# mailto:tommy@tmtm.org

require "socket"
require "timeout"
require "digest/sha1"
require "stringio"

class RbMysql
  # MySQL network protocol
  class Protocol

    VERSION = 10
    MAX_PACKET_LENGTH = 2**24-1

    # Convert netdata to Ruby value
    # === Argument
    # data :: [Packet] packet data
    # type :: [Integer] field type
    # unsigned :: [true or false] true if value is unsigned
    # === Return
    # Object :: converted value.
    def self.net2value(pkt, type, unsigned)
      case type
      when Field::TYPE_STRING, Field::TYPE_VAR_STRING, Field::TYPE_NEWDECIMAL, Field::TYPE_BLOB, Field::TYPE_JSON
        return pkt.lcs
      when Field::TYPE_TINY
        v = pkt.utiny
        return unsigned ? v : v < 128 ? v : v-256
      when Field::TYPE_SHORT
        v = pkt.ushort
        return unsigned ? v : v < 32768 ? v : v-65536
      when Field::TYPE_INT24, Field::TYPE_LONG
        v = pkt.ulong
        return unsigned ? v : v < 0x8000_0000 ? v : v-0x10000_0000
      when Field::TYPE_LONGLONG
        n1, n2 = pkt.ulong, pkt.ulong
        v = (n2 << 32) | n1
        return unsigned ? v : v < 0x8000_0000_0000_0000 ? v : v-0x10000_0000_0000_0000
      when Field::TYPE_FLOAT
        return pkt.read(4).unpack('e').first
      when Field::TYPE_DOUBLE
        return pkt.read(8).unpack('E').first
      when Field::TYPE_DATE
        len = pkt.utiny
        y, m, d = pkt.read(len).unpack("vCC")
        t = RbMysql::Time.new(y, m, d, nil, nil, nil)
        return t
      when Field::TYPE_DATETIME, Field::TYPE_TIMESTAMP
        len = pkt.utiny
        y, m, d, h, mi, s, sp = pkt.read(len).unpack("vCCCCCV")
        return RbMysql::Time.new(y, m, d, h, mi, s, false, sp)
      when Field::TYPE_TIME
        len = pkt.utiny
        sign, d, h, mi, s, sp = pkt.read(len).unpack("CVCCCV")
        h = d.to_i * 24 + h.to_i
        return RbMysql::Time.new(0, 0, 0, h, mi, s, sign!=0, sp)
      when Field::TYPE_YEAR
        return pkt.ushort
      when Field::TYPE_BIT
        return pkt.lcs
      else
        raise "not implemented: type=#{type}"
      end
    end

    # convert Ruby value to netdata
    # === Argument
    # v :: [Object] Ruby value.
    # === Return
    # Integer :: type of column. Field::TYPE_*
    # String :: netdata
    # === Exception
    # ProtocolError :: value too large / value is not supported
    def self.value2net(v)
      case v
      when nil
        type = Field::TYPE_NULL
        val = ""
      when Integer
        if -0x8000_0000 <= v && v < 0x8000_0000
          type = Field::TYPE_LONG
          val = [v].pack('V')
        elsif -0x8000_0000_0000_0000 <= v && v < 0x8000_0000_0000_0000
          type = Field::TYPE_LONGLONG
          val = [v&0xffffffff, v>>32].pack("VV")
        elsif 0x8000_0000_0000_0000 <= v && v <= 0xffff_ffff_ffff_ffff
          type = Field::TYPE_LONGLONG | 0x8000
          val = [v&0xffffffff, v>>32].pack("VV")
        else
          raise ProtocolError, "value too large: #{v}"
        end
      when Float
        type = Field::TYPE_DOUBLE
        val = [v].pack("E")
      when String
        type = Field::TYPE_STRING
        val = Packet.lcs(v)
      when ::Time
        type = Field::TYPE_DATETIME
        val = [11, v.year, v.month, v.day, v.hour, v.min, v.sec, v.usec].pack("CvCCCCCV")
      when RbMysql::Time
        type = Field::TYPE_DATETIME
        val = [11, v.year, v.month, v.day, v.hour, v.min, v.sec, v.second_part].pack("CvCCCCCV")
      else
        raise ProtocolError, "class #{v.class} is not supported"
      end
      return type, val
    end

    attr_reader :server_info
    attr_reader :server_version
    attr_reader :thread_id
    attr_reader :sqlstate
    attr_reader :affected_rows
    attr_reader :insert_id
    attr_reader :server_status
    attr_reader :warning_count
    attr_reader :message
    attr_accessor :charset

    # @state variable keep state for connection.
    # :INIT   :: Initial state.
    # :READY  :: Ready for command.
    # :FIELD  :: After query(). retr_fields() is needed.
    # :RESULT :: After retr_fields(), retr_all_records() or stmt_retr_all_records() is needed.

    # make socket connection to server.
    # === Argument
    # host :: [String] if "localhost" or "" nil then use UNIXSocket. Otherwise use TCPSocket
    # port :: [Integer] port number using by TCPSocket
    # socket :: [String,Socket] socket file name using by UNIXSocket, or an existing ::Socket instance
    # conn_timeout :: [Integer] connect timeout (sec).
    # read_timeout :: [Integer] read timeout (sec).
    # write_timeout :: [Integer] write timeout (sec).
    # === Exception
    # [ClientError] :: connection timeout
    def initialize(host, port, socket, conn_timeout, read_timeout, write_timeout)
      @insert_id = 0
      @warning_count = 0
      @gc_stmt_queue = []   # stmt id list which GC destroy.
      set_state :INIT
      @read_timeout = read_timeout
      @write_timeout = write_timeout
      begin
        Timeout.timeout conn_timeout do
          if host.nil? or host.empty? or host == "localhost"
            socket ||= ENV["MYSQL_UNIX_PORT"] || MYSQL_UNIX_PORT
            @sock = UNIXSocket.new socket
          else
            if !socket
              port ||= ENV["MYSQL_TCP_PORT"] || (Socket.getservbyname("mysql","tcp") rescue MYSQL_TCP_PORT)
              @sock = TCPSocket.new host, port
            else
              @sock = socket
            end
          end
        end
      rescue Timeout::Error
        raise ClientError, "connection timeout"
      end
    end

    def close
      @sock.close
    end

    # initial negotiate and authenticate.
    # === Argument
    # user    :: [String / nil] username
    # passwd  :: [String / nil] password
    # db      :: [String / nil] default database name. nil: no default.
    # flag    :: [Integer] client flag
    # charset :: [RbMysql::Charset / nil] charset for connection. nil: use server's charset
    # === Exception
    # ProtocolError :: The old style password is not supported
    def authenticate(user, passwd, db, flag, charset)
      check_state :INIT
      @authinfo = [user, passwd, db, flag, charset]
      reset
      init_packet = InitialPacket.parse read
      @server_info = init_packet.server_version
      @server_version = init_packet.server_version.split(/\D/)[0,3].inject{|a,b|a.to_i*100+b.to_i}
      @thread_id = init_packet.thread_id
      @client_flags = CLIENT_LONG_PASSWORD | CLIENT_LONG_FLAG | CLIENT_TRANSACTIONS | CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION
      @client_flags |= CLIENT_CONNECT_WITH_DB if db
      @client_flags |= flag
      @charset = charset
      unless @charset
        @charset = Charset.by_number(init_packet.server_charset)
        @charset.encoding       # raise error if unsupported charset
      end
      netpw = encrypt_password passwd, init_packet.scramble_buff
      write AuthenticationPacket.serialize(@client_flags, 1024**3, @charset.number, user, netpw, db)
      raise ProtocolError, 'The old style password is not supported' if read.to_s == "\xfe"
      set_state :READY
    end

    # Quit command
    def quit_command
      synchronize do
        reset
        write [COM_QUIT].pack("C")
        close
      end
    end

    # Query command
    # === Argument
    # query :: [String] query string
    # === Return
    # [Integer / nil] number of fields of results. nil if no results.
    def query_command(query)
      check_state :READY
      begin
        reset
        write [COM_QUERY, @charset.convert(query)].pack("Ca*")
        get_result
      rescue
        set_state :READY
        raise
      end
    end

    # get result of query.
    # === Return
    # [integer / nil] number of fields of results. nil if no results.
    def get_result
      begin
        res_packet = ResultPacket.parse read
        if res_packet.field_count.to_i > 0  # result data exists
          set_state :FIELD
          return res_packet.field_count
        end
        if res_packet.field_count.nil?      # LOAD DATA LOCAL INFILE
          if @client_flags.to_i & CLIENT_LOCAL_FILES == 0
            raise ProtocolError, 'Load data local infile forbidden'
          end
          filename = res_packet.message
          File.open(filename){|f| write f}
          write nil  # EOF mark
          read
        end
        @affected_rows, @insert_id, @server_status, @warning_count, @message =
          res_packet.affected_rows, res_packet.insert_id, res_packet.server_status, res_packet.warning_count, res_packet.message
        set_state :READY
        return nil
      rescue
        set_state :READY
        raise
      end
    end

    # Retrieve n fields
    # === Argument
    # n :: [Integer] number of fields
    # === Return
    # [Array of RbMysql::Field] field list
    def retr_fields(n)
      check_state :FIELD
      begin
        fields = n.times.map{Field.new FieldPacket.parse(read)}
        read_eof_packet
        set_state :RESULT
        fields
      rescue
        set_state :READY
        raise
      end
    end

    # Retrieve all records for simple query
    # === Argument
    # fields :: [Array<RbMysql::Field>] number of fields
    # === Return
    # [Array of Array of String] all records
    def retr_all_records(fields)
      check_state :RESULT
      enc = charset.encoding
      begin
        all_recs = []
        until (pkt = read).eof?
          all_recs.push RawRecord.new(pkt, fields, enc)
        end
        pkt.read(3)
        @server_status = pkt.utiny
        all_recs
      ensure
        set_state :READY
      end
    end

    # Field list command
    # === Argument
    # table :: [String] table name.
    # field :: [String / nil] field name that may contain wild card.
    # === Return
    # [Array of Field] field list
    def field_list_command(table, field)
      synchronize do
        reset
        write [COM_FIELD_LIST, table, 0, field].pack("Ca*Ca*")
        fields = []
        until (data = read).eof?
          fields.push Field.new(FieldPacket.parse(data))
        end
        return fields
      end
    end

    # Process info command
    # === Return
    # [Array of Field] field list
    def process_info_command
      check_state :READY
      begin
        reset
        write [COM_PROCESS_INFO].pack("C")
        field_count = read.lcb
        fields = field_count.times.map{Field.new FieldPacket.parse(read)}
        read_eof_packet
        set_state :RESULT
        return fields
      rescue
        set_state :READY
        raise
      end
    end

    # Ping command
    def ping_command
      simple_command [COM_PING].pack("C")
    end

    # Kill command
    def kill_command(pid)
      simple_command [COM_PROCESS_KILL, pid].pack("CV")
    end

    # Refresh command
    def refresh_command(op)
      simple_command [COM_REFRESH, op].pack("CC")
    end

    # Set option command
    def set_option_command(opt)
      simple_command [COM_SET_OPTION, opt].pack("Cv")
    end

    # Shutdown command
    def shutdown_command(level)
      simple_command [COM_SHUTDOWN, level].pack("CC")
    end

    # Statistics command
    def statistics_command
      simple_command [COM_STATISTICS].pack("C")
    end

    # Stmt prepare command
    # === Argument
    # stmt :: [String] prepared statement
    # === Return
    # [Integer] statement id
    # [Integer] number of parameters
    # [Array of Field] field list
    def stmt_prepare_command(stmt)
      synchronize do
        reset
        write [COM_STMT_PREPARE, charset.convert(stmt)].pack("Ca*")
        res_packet = PrepareResultPacket.parse read
        if res_packet.param_count > 0
          res_packet.param_count.times{read}    # skip parameter packet
          read_eof_packet
        end
        if res_packet.field_count > 0
          fields = res_packet.field_count.times.map{Field.new FieldPacket.parse(read)}
          read_eof_packet
        else
          fields = []
        end
        return res_packet.statement_id, res_packet.param_count, fields
      end
    end

    # Stmt execute command
    # === Argument
    # stmt_id :: [Integer] statement id
    # values  :: [Array] parameters
    # === Return
    # [Integer] number of fields
    def stmt_execute_command(stmt_id, values)
      check_state :READY
      begin
        reset
        write ExecutePacket.serialize(stmt_id, RbMysql::Stmt::CURSOR_TYPE_NO_CURSOR, values)
        get_result
      rescue
        set_state :READY
        raise
      end
    end

    # Retrieve all records for prepared statement
    # === Argument
    # fields  :: [Array of RbMysql::Fields] field list
    # charset :: [RbMysql::Charset]
    # === Return
    # [Array of Array of Object] all records
    def stmt_retr_all_records(fields, charset)
      check_state :RESULT
      enc = charset.encoding
      begin
        all_recs = []
        until (pkt = read).eof?
          all_recs.push StmtRawRecord.new(pkt, fields, enc)
        end
        all_recs
      ensure
        set_state :READY
      end
    end

    # Stmt close command
    # === Argument
    # stmt_id :: [Integer] statement id
    def stmt_close_command(stmt_id)
      synchronize do
        reset
        write [COM_STMT_CLOSE, stmt_id].pack("CV")
      end
    end

    def gc_stmt(stmt_id)
      @gc_stmt_queue.push stmt_id
    end

    private

    def check_state(st)
      raise 'command out of sync' unless @state == st
    end

    def set_state(st)
      @state = st
      if st == :READY
        gc_disabled = GC.disable
        begin
          while st = @gc_stmt_queue.shift
            reset
            write [COM_STMT_CLOSE, st].pack("CV")
          end
        ensure
          GC.enable unless gc_disabled
        end
      end
    end

    def synchronize
      begin
        check_state :READY
        return yield
      ensure
        set_state :READY
      end
    end

    # Reset sequence number
    def reset
      @seq = 0    # packet counter. reset by each command
    end

    # Read one packet data
    # === Return
    # [Packet] packet data
    # === Exception
    # [ProtocolError] invalid packet sequence number
    def read
      data = ''
      len = nil
      begin
        Timeout.timeout @read_timeout do
          header = @sock.read(4)
          raise EOFError unless header && header.length == 4
          len1, len2, seq = header.unpack("CvC")
          len = (len2 << 8) + len1
          raise ProtocolError, "invalid packet: sequence number mismatch(#{seq} != #{@seq}(expected))" if @seq != seq
          @seq = (@seq + 1) % 256
          ret = @sock.read(len)
          raise EOFError unless ret && ret.length == len
          data.concat ret
        end
      rescue EOFError
        raise ClientError::ServerGoneError, 'MySQL server has gone away'
      rescue Timeout::Error
        raise ClientError, "read timeout"
      end while len == MAX_PACKET_LENGTH

      @sqlstate = "00000"

      # Error packet
      if data[0] == ?\xff
        f, errno, marker, @sqlstate, message = data.unpack("Cvaa5a*")
        unless marker == "#"
          f, errno, message = data.unpack("Cva*")    # Version 4.0 Error
          @sqlstate = ""
        end
        message.force_encoding(@charset.encoding) if @charset
        if RbMysql::ServerError::ERROR_MAP.key? errno
          raise RbMysql::ServerError::ERROR_MAP[errno].new(message, @sqlstate)
        end
        raise RbMysql::ServerError.new(message, @sqlstate)
      end
      Packet.new(data)
    end

    # Write one packet data
    # === Argument
    # data :: [String / IO] packet data. If data is nil, write empty packet.
    def write(data)
      begin
        @sock.sync = false
        if data.nil?
          Timeout.timeout @write_timeout do
            @sock.write [0, 0, @seq].pack("CvC")
          end
          @seq = (@seq + 1) % 256
        else
          data = StringIO.new data if data.is_a? String
          while d = data.read(MAX_PACKET_LENGTH)
            Timeout.timeout @write_timeout do
              @sock.write [d.length%256, d.length/256, @seq].pack("CvC")
              @sock.write d
            end
            @seq = (@seq + 1) % 256
          end
        end
        @sock.sync = true
        Timeout.timeout @write_timeout do
          @sock.flush
        end
      rescue Errno::EPIPE
        raise ClientError::ServerGoneError, 'MySQL server has gone away'
      rescue Timeout::Error
        raise ClientError, "write timeout"
      end
    end

    # Read EOF packet
    # === Exception
    # [ProtocolError] packet is not EOF
    def read_eof_packet
      raise ProtocolError, "packet is not EOF" unless read.eof?
    end

    # Send simple command
    # === Argument
    # packet :: [String] packet data
    # === Return
    # [String] received data
    def simple_command(packet)
      synchronize do
        reset
        write packet
        read.to_s
      end
    end

    # Encrypt password
    # === Argument
    # plain    :: [String] plain password.
    # scramble :: [String] scramble code from initial packet.
    # === Return
    # [String] encrypted password
    def encrypt_password(plain, scramble)
      return "" if plain.nil? or plain.empty?
      hash_stage1 = Digest::SHA1.digest plain
      hash_stage2 = Digest::SHA1.digest hash_stage1
      return hash_stage1.unpack("C*").zip(Digest::SHA1.digest(scramble+hash_stage2).unpack("C*")).map{|a,b| a^b}.pack("C*")
    end

    # Initial packet
    class InitialPacket
      def self.parse(pkt)
        protocol_version = pkt.utiny
        server_version = pkt.string
        thread_id = pkt.ulong
        scramble_buff = pkt.read(8)
        f0 = pkt.utiny
        server_capabilities = pkt.ushort
        server_charset = pkt.utiny
        server_status = pkt.ushort
        _f1 = pkt.read(13)
        rest_scramble_buff = pkt.string
        raise ProtocolError, "unsupported version: #{protocol_version}" unless protocol_version == VERSION
        raise ProtocolError, "invalid packet: f0=#{f0}" unless f0 == 0
        scramble_buff.concat rest_scramble_buff
        self.new protocol_version, server_version, thread_id, server_capabilities, server_charset, server_status, scramble_buff
      end

      attr_reader :protocol_version, :server_version, :thread_id, :server_capabilities, :server_charset, :server_status, :scramble_buff

      def initialize(*args)
        @protocol_version, @server_version, @thread_id, @server_capabilities, @server_charset, @server_status, @scramble_buff = args
      end
    end

    # Result packet
    class ResultPacket
      def self.parse(pkt)
        field_count = pkt.lcb
        if field_count == 0
          affected_rows = pkt.lcb
          insert_id = pkt.lcb
          server_status = pkt.ushort
          warning_count = pkt.ushort
          message = pkt.lcs
          return self.new(field_count, affected_rows, insert_id, server_status, warning_count, message)
        elsif field_count.nil?   # LOAD DATA LOCAL INFILE
          return self.new(nil, nil, nil, nil, nil, pkt.to_s)
        else
          return self.new(field_count)
        end
      end

      attr_reader :field_count, :affected_rows, :insert_id, :server_status, :warning_count, :message

      def initialize(*args)
        @field_count, @affected_rows, @insert_id, @server_status, @warning_count, @message = args
      end
    end

    # Field packet
    class FieldPacket
      def self.parse(pkt)
        _first = pkt.lcs
        db = pkt.lcs
        table = pkt.lcs
        org_table = pkt.lcs
        name = pkt.lcs
        org_name = pkt.lcs
        _f0 = pkt.utiny
        charsetnr = pkt.ushort
        length = pkt.ulong
        type = pkt.utiny
        flags = pkt.ushort
        decimals = pkt.utiny
        f1 = pkt.ushort

        raise ProtocolError, "invalid packet: f1=#{f1}" unless f1 == 0
        default = pkt.lcs
        return self.new(db, table, org_table, name, org_name, charsetnr, length, type, flags, decimals, default)
      end

      attr_reader :db, :table, :org_table, :name, :org_name, :charsetnr, :length, :type, :flags, :decimals, :default

      def initialize(*args)
        @db, @table, @org_table, @name, @org_name, @charsetnr, @length, @type, @flags, @decimals, @default = args
      end
    end

    # Prepare result packet
    class PrepareResultPacket
      def self.parse(pkt)
        raise ProtocolError, "invalid packet" unless pkt.utiny == 0
        statement_id = pkt.ulong
        field_count = pkt.ushort
        param_count = pkt.ushort
        f = pkt.utiny
        warning_count = pkt.ushort
        raise ProtocolError, "invalid packet" unless f == 0x00
        self.new statement_id, field_count, param_count, warning_count
      end

      attr_reader :statement_id, :field_count, :param_count, :warning_count

      def initialize(*args)
        @statement_id, @field_count, @param_count, @warning_count = args
      end
    end

    # Authentication packet
    class AuthenticationPacket
      def self.serialize(client_flags, max_packet_size, charset_number, username, scrambled_password, databasename)
        [
          client_flags,
          max_packet_size,
          Packet.lcb(charset_number),
          "",                   # always 0x00 * 23
          username,
          Packet.lcs(scrambled_password),
          databasename
        ].pack("VVa*a23Z*A*Z*")
      end
    end

    # Execute packet
    class ExecutePacket
      def self.serialize(statement_id, cursor_type, values)
        nbm = null_bitmap values
        netvalues = ""
        types = values.map do |v|
          t, n = Protocol.value2net v
          netvalues.concat n if v
          t
        end
        [RbMysql::COM_STMT_EXECUTE, statement_id, cursor_type, 1, nbm, 1, types.pack("v*"), netvalues].pack("CVCVa*Ca*a*")
      end

      # make null bitmap
      #
      # If values is [1, nil, 2, 3, nil] then returns "\x12"(0b10010).
      def self.null_bitmap(values)
        bitmap = values.enum_for(:each_slice,8).map do |vals|
          vals.reverse.inject(0){|b, v|(b << 1 | (v ? 0 : 1))}
        end
        return bitmap.pack("C*")
      end

    end
  end

  class RawRecord
    def initialize(packet, fields, encoding)
      @packet, @fields, @encoding = packet, fields, encoding
    end

    def to_a
      @fields.map do |f|
        if s = @packet.lcs
          unless f.type == Field::TYPE_BIT or f.charsetnr == Charset::BINARY_CHARSET_NUMBER
            s = Charset.convert_encoding(s, @encoding)
          end
        end
        s
      end
    end
  end

  class StmtRawRecord
    # === Argument
    # pkt     :: [Packet]
    # fields  :: [Array of Fields]
    # encoding:: [Encoding]
    def initialize(packet, fields, encoding)
      @packet, @fields, @encoding = packet, fields, encoding
    end

    # Parse statement result packet
    # === Return
    # [Array of Object] one record
    def parse_record_packet
      @packet.utiny  # skip first byte
      null_bit_map = @packet.read((@fields.length+7+2)/8).unpack("b*").first
      rec = @fields.each_with_index.map do |f, i|
        if null_bit_map[i+2] == ?1
          nil
        else
          unsigned = f.flags & Field::UNSIGNED_FLAG != 0
          v = Protocol.net2value(@packet, f.type, unsigned)
          if v.is_a? Numeric or v.is_a? RbMysql::Time
            v
          elsif f.type == Field::TYPE_BIT or f.charsetnr == Charset::BINARY_CHARSET_NUMBER
            Charset.to_binary(v)
          else
            Charset.convert_encoding(v, @encoding)
          end
        end
      end
      rec
    end

    alias to_a parse_record_packet

  end
end