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', 0] )
			], 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)
		otimeout = dcerpc.options['read_timeout']
		
		begin
			dcerpc.options['read_timeout'] = timeout if timeout
			dcerpc.call(function, stub)
		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