require 'msf/core'

module Msf

###
#
# ExploitDriver
# -------------
#
# This class drives the exploitation process from start to finish for a given
# exploit module instance.  It's responsible for payload generation, encoding,
# and padding as well as initialization handlers and finally launching the
# exploit.
#
###
class ExploitDriver

	def initialize(framework)
		self.payload    = nil
		self.exploit    = nil
		self.target_idx = nil
	end

	#
	# Specification of the exploit target index
	#
	def target_idx=(target_idx)
		if (target_idx)
			# Make sure the target index is valid
			if (target_idx >= exploit.targets.length)
				raise Rex::ArgumentError, "Invalid target index.", caller
			end
		end

		 # Set the active target
		@target_idx = target_idx
	end

	#
	# Checks to see if the supplied payload is compatible with the
	# current exploit.  Assumes that target_idx is valid.
	#
	def compatible_payload?(payload)
		return ((payload.platform & exploit.targets[target_idx].platform).empty? == false)
	end

	##
	#
	# Exploit execution
	#
	##

	#
	# Makes sure everything's in tip-top condition prior to launching the
	# exploit.  For things that aren't good to go, an exception is thrown.
	#
	def validate
		# First, validate that a target has been selected
		if (target_idx == nil)
			raise MissingTargetError, 
				"A payload cannot be selected until a target is specified.",
				caller
		end

		# Next, validate that a payload has been selected
		if (payload == nil)
			raise MissingPayloadError, 
				"A payload has not been selected.", caller
		end

		# Make sure the payload is compatible after all
		if (compatible_payload?(payload) == false)
			raise IncompatiblePayloadError.new(payload.refname), 
				"Incompatible payload", caller
		end

		# Finally, validate options on the exploit module to ensure that things
		# are ready to operate as they should.
		exploit.options.validate(exploit.datastore)

		return true
	end

	#
	# Kicks off an exploitation attempt and performs the following four major
	# operations:
	#
	#   - Generates the payload
	#   - Initializes & monitors the handler
	#   - Launches the exploit
	#   - Cleans up the handler
	#
	def run
		# First thing's first -- validate the state.  Make sure all requirement
		# parameters are set, including those that are derived from the
		# datastore.
		validate()
	
		# After validation has occurred, it's time to set some values on the
		# exploit instance and begin preparing the payload
		exploit.datastore['TARGET'] = @target_idx
	
		# Generate the encoded version of the supplied payload on the exploit
		# module instance
		exploit.generate_payload(payload)

		# Default the session to nil
		session = nil

		begin
			# Set the exploit up the bomb
			exploit.setup

			# Launch the exploit
			exploit.exploit

			# Wait the payload to acquire a session
			session = payload.wait_for_session
		ensure
			# Ensure that, no matter what, clean up of the handler occurs
			payload.stop_handler
		
			# Allow the exploit to cleanup after itself, that messy bugger.
			exploit.cleanup
		end

		# Set the vector through which this session was created
		if (session)
			session.set_via(
				'Exploit' => exploit.refname,
				'Payload' => payload.refname)
		end

		return session
	end

	attr_reader   :target_idx
	attr_accessor :exploit
	attr_accessor :payload

end

end