require 'msf/core'
require 'xdr'

module Msf

###
#
# This module exposes methods for accessing NDMP services using a class-based wrapper
# around normal sockets, allowing the use of more than one NDMP connection at once.
#
###
module Exploit::Remote::NDMPSocket

module NDMP

#
# This class represents a NDMP message, including its header and body.
#
class Message
  class Header < XDR::Struct
    attribute :sequence_num,       XDR::Int
    attribute :timestamp,          XDR::Int
    attribute :is_response,        XDR::Bool
    attribute :type,               XDR::Int
    attribute :reply_sequence_num, XDR::Int
    attribute :error,              XDR::Int
  end

  CONFIG_GET_HOST_INFO   = 0x100
  CONFIG_GET_BUTYPE_ATTR = 0x101
  CONFIG_GET_SERVER_INFO = 0x108
  NOTIFY_CONNECTED       = 0x502
  CONNECT_OPEN           = 0x900
  CONNECT_CLIENT_AUTH    = 0x901

  def self.new_request(type, body='')
    header = Header.new(
      :sequence_num       => nil,
      :timestamp          => nil,
      :is_response        => false,
      :type               => type,
      :reply_sequence_num => 0,
      :error              => 0
    )
    new(header, body)
  end

  attr_accessor :header, :body

  def initialize(header, body)
    @header = header
    @body = body
  end
end

#
# This class wraps a normal socket with NDMP functionality, such as NDMP message reading
# and writing.
#
class Socket
  def initialize(sock)
    @sock = sock
    @next_sequence_num = 1
  end

  def raw_recv(*args)
    @sock.recv(*args)
  end

  def raw_recvall(n, *args)
    r = ''
    while r.length < n
      s = raw_recv(n - r.length, *args)
      return nil if s.to_s.empty?
      r << s
    end
    r
  end

  def raw_send(*args)
    @sock.send(*args)
  end

  def raw_sendall(s, *args)
    n = 0
    while n < s.length
      i = raw_send(s[n..s.length], *args)
      return false if i <= 0
      n += i
    end

    true
  end

  def close
    @sock.close
  end

  #
  # Read a single NDMP message. If require_ok_type is given, the message must be of the
  # given type and have no error indicated.
  #
  def read_ndmp_msg(require_ok_type=nil)
    frags = read_ndmp_frags
    return nil if frags.nil?
    header = Message::Header.from_xdr(frags.slice!(0...(4 * 6)))

    return nil if require_ok_type && (require_ok_type != header.type || header.error != 0)

    Message.new(header, frags)
  end

  #
  # Prepare a NDMP message for sending and send it. Can send messages multiple times.
  #
  # If all_but_all_char is true, then the last character will be held back and will be
  # returned so that it can be sent at a later point elsewhere. This is sometimes
  # necessary for exploiting e.g. race conditions.
  #
  def prepare_and_write_ndmp_msg(msg, all_but_last_char=false, times=1, flags=0)
    msg.header.sequence_num = @next_sequence_num
    @next_sequence_num += 1
    msg.header.timestamp = Time.now.to_i

    frag = msg.header.to_xdr + msg.body
    write_ndmp_frag(frag, all_but_last_char, times, flags)
  end

  #
  # Send and recieve a pair of NDMP messages.
  #
  def do_request_response(msg, *args)
    return nil unless prepare_and_write_ndmp_msg(msg, *args)
    read_ndmp_msg(msg.header.type)
  end

  #
  # Establish a SSL session on the socket. Raw socket reading/writing functions are
  # replaced with their SSL equivalents.
  #
  def wrap_with_ssl(ssl_context)
    @sock = OpenSSL::SSL::SSLSocket.new(@sock, ssl_context)
    @sock.connect

    def self.raw_recv(n, *_args)
      @sock.sysread(n)
    end

    def self.raw_send(b, *_args)
      @sock.syswrite(b)
    end
  end

  private

  #
  # Read and reassemble a group of NDMP fragments. Usually NDMP messages are sent in a
  # single NDMP fragment, but this is not guaranteed by the standard.
  #
  def read_ndmp_frags
    result = ''

    loop do
      buf = raw_recvall(4)
      return nil if buf.nil?
      n = buf.unpack('N')[0]
      len = n & 0x7fffffff
      last = (n & 0x80000000) != 0

      buf = raw_recvall(len)
      return nil if buf.nil?
      result << buf

      break if last
    end

    result
  end

  #
  # Write a NDMP fragment (i.e. containing a NDMP packet).
  #
  # Can hold back on sending the last character; see prepare_and_write_ndmp_msg.
  #
  def write_ndmp_frag(buf, all_but_last_char, times, flags)
    buf = ([buf.length | 0x80000000].pack('N') + buf) * times

    return false unless raw_sendall(all_but_last_char ? buf[0...-1] : buf, flags)

    all_but_last_char ? buf[-1] : true
  end

end

end

end

end