require 'rex/proto/dcerpc'
require 'rex/encoder/ndr'

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
	include Exploit::Remote::Tcp

	# 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

	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
		)
	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

		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 self.respond_to?('simple') and self.simple
			opts['smb_client'] = self.simple
		end

		if (self.sock and not self.dcerpc_socket)
			self.dcerpc_socket = self.sock
		end

		self.dcerpc = Rex::Proto::DCERPC::Client.new(h, self.dcerpc_socket, 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 = '')
		begin
			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
		rescue ::Exception
			raise $!
#			print_status("DCERPC Error: #{$!.class.to_s} - #{$!.backtrace.join("  -  ")}")
#			return
		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, :dcerpc_socket
	
end

end