# -*- coding: binary -*-
require 'rex/proto/dcerpc'
require 'rex/encoder/ndr'
require 'msf/core/exploit/dcerpc_epm'
require 'msf/core/exploit/dcerpc_mgmt'
require 'msf/core/exploit/dcerpc_lsa'

module Msf

###
#
# This mixin provides utility methods for interacting with a DCERPC service on
# a remote machine.  These methods may generally be useful in the context of
# exploitation.  This mixin extends the Tcp exploit mixin. Only one DCERPC
# service can be accessed at a time using this class.
#
###
module Exploit::Remote::DCERPC

  # Alias over the Rex DCERPC protocol modules
  DCERPCPacket   = Rex::Proto::DCERPC::Packet
  DCERPCClient   = Rex::Proto::DCERPC::Client
  DCERPCResponse = Rex::Proto::DCERPC::Response
  DCERPCUUID     = Rex::Proto::DCERPC::UUID
  NDR            = Rex::Encoder::NDR


  # Support TCP-based RPC services
  include Exploit::Remote::Tcp

  # Helper methods for specific services
  include Exploit::Remote::DCERPC_EPM
  include Exploit::Remote::DCERPC_MGMT
  include Exploit::Remote::DCERPC_LSA

  def initialize(info = {})
    super

    register_evasion_options(
      [
        OptInt.new('DCERPC::max_frag_size',   [ true, 'Set the DCERPC packet fragmentation size', 4096]),
        OptBool.new('DCERPC::fake_bind_multi', [ false, 'Use multi-context bind calls', true ]),
        OptInt.new('DCERPC::fake_bind_multi_prepend',  [ false, 'Set the number of UUIDs to prepend before the target', 0]),
        OptInt.new('DCERPC::fake_bind_multi_append',   [ false, 'Set the number of UUIDs to append the target', 0]),
        OptEnum.new('DCERPC::smb_pipeio', [ false, 'Use a different delivery method for accessing named pipes', 'rw', ['rw', 'trans']] )

      ], Msf::Exploit::Remote::DCERPC)

    register_options(
      [
        Opt::RHOST,
        Opt::RPORT(135),
      ], Msf::Exploit::Remote::DCERPC)

    register_advanced_options(
      [
        OptInt.new('DCERPC::ReadTimeout',   [ true, 'The number of seconds to wait for DCERPC responses', 10] )
      ], Msf::Exploit::Remote::DCERPC)

  end

  def dcerpc_handle(uuid, version, protocol, opts)
    self.handle = Rex::Proto::DCERPC::Handle.new([uuid, version], protocol, rhost, opts)
  end

  def dcerpc_bind(h)
    opts = { 'Msf' => framework, 'MsfExploit' => self }

    if datastore['DCERPC::max_frag_size']
      opts['frag_size'] = datastore['DCERPC::max_frag_size']
    end

    if datastore['DCERPC::fake_bind_multi']
      opts['fake_multi_bind'] = 1

      if datastore['DCERPC::fake_bind_multi_prepend']
        opts['fake_multi_bind_prepend'] = datastore['DCERPC::fake_bind_multi_prepend']
      end

      if datastore['DCERPC::fake_bind_multi_append']
        opts['fake_multi_bind_append'] = datastore['DCERPC::fake_bind_multi_append']
      end
    end

    opts['connect_timeout'] = (datastore['ConnectTimeout'] || 10).to_i

    opts['read_timeout']    = (datastore['DCERPC::ReadTimeout'] || 10).to_i


    # Configure the SMB evasion options

    if (datastore['SMBUser'])
      opts['smb_user'] = datastore['SMBUser']
    end

    if (datastore['SMBPass'])
      opts['smb_pass'] = datastore['SMBPass']
    end

    if (datastore['DCERPC::smb_pipeio'])
      opts['smb_pipeio'] = datastore['DCERPC::smb_pipeio']
    end

    if (datastore['SMB::pipe_write_min_size'])
      opts['pipe_write_min_size'] = datastore['SMB::pipe_write_min_size']
    end

    if (datastore['SMB::pipe_write_max_size'])
      opts['pipe_write_max_size'] = datastore['SMB::pipe_write_max_size']
    end

    if (datastore['SMB::pipe_read_min_size'])
      opts['pipe_read_min_size'] = datastore['SMB::pipe_read_min_size']
    end

    if (datastore['SMB::pipe_read_max_size'])
      opts['pipe_read_max_size'] = datastore['SMB::pipe_read_max_size']
    end

    if (self.respond_to?('simple') and self.simple)
      opts['smb_client'] = self.simple
    end

    # Create the DCERPC client
    self.dcerpc = Rex::Proto::DCERPC::Client.new(h, self.sock, opts)

    if (self.handle.protocol == 'ncacn_np' and not self.simple)
      self.simple = self.dcerpc.smb  # expose the simple client if we have access to it
    end
  end

  def dcerpc_call(function, stub = '', timeout=nil, do_recv=true)
    otimeout = dcerpc.options['read_timeout']

    begin
      dcerpc.options['read_timeout'] = timeout if timeout
      dcerpc.call(function, stub, do_recv)
    rescue ::Rex::Proto::SMB::Exceptions::NoReply, Rex::Proto::DCERPC::Exceptions::NoResponse
      print_status("The DCERPC service did not reply to our request")
      return
    ensure
      dcerpc.options['read_timeout'] = otimeout
    end
  end

  # Convert a standard ASCII string to 16-bit Unicode
  def unicode(str)
    Rex::Text.to_unicode(str)
  end

  # Useful accessors for tracking DCERPC state
  attr_accessor	:handle, :dcerpc

end

end