require 'rex/proto/dcerpc'

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::Proto::DCERPC::NDR

	def initialize(info = {})
		super
	
		register_advanced_options(
			[
				OptInt.new('DCERPCFragSize',   [ true, 'Set the DCERPC packet fragmentation size', 127]),
				OptBool.new('DCERPCFakeMultiBind', [ false, 'Use multi-context bind calls', 'True' ])
			], 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, datastore['RHOST'], opts)
	end

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

		if datastore['DCERPCFragSize']
			opts['frag_size'] = datastore['DCERPCFragSize']
		end

		if datastore['DCERPCFakeMultiBind']
			opts['fake_multi_bind'] = 1
		end

		if datastore['SMBUSER']
			opts['smb_user'] = datastore['SMBUSER']
		end
		
		if datastore['SMBPASS']
			opts['smb_pass'] = datastore['SMBPASS']
		end

		self.dcerpc = Rex::Proto::DCERPC::Client.new(h, dcerpc_socket(), opts)

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

	def dcerpc_call (function, stub = '')
		dcerpc.call(function, stub)
	end

	# Convert a standard ASCII string to 16-bit Unicode
	def unicode (str)
		Rex::Text.to_unicode(str)
	end
	
	# Used to track the last DCERPC context 
	attr_accessor	:dcerpc_bind_context, :handle, :dcerpc, :dcerpc_socket
	
end

end