require 'msf/core'

module Msf

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

	#
	# Initializes the exploit driver using the supplied framework instance.
	#
	def initialize(framework)
		self.payload                = nil
		self.exploit                = nil
		self.use_job                = false
		self.force_wait_for_session = false
	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

	#
	# This method returns the currently selected target index.
	#
	def 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

		# Associate the payload instance with the exploit
		payload.assoc_exploit = exploit

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

		# Validate the payload's options.  The payload's datastore is
		# most likely shared against the exploit's datastore, but in case it
		# isn't.
		payload.options.validate(payload.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
		self.session = nil

		# Set up the run context
		ctx = [ exploit, payload ]

		# If we are being instructed to run as a job then let's create that job
		# like a good person.
		if (use_job)
			exploit.framework.jobs.start_bg_job(
				"Exploit: #{exploit.refname}", 
				ctx,
				Proc.new { |ctx| job_run_proc(ctx) },
				Proc.new { |ctx| job_cleanup_proc(ctx) })
		# Otherwise, run inline.
		else
			job_run_proc(ctx)
			job_cleanup_proc(ctx)
		end

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

			# Call the exploit success handler
			exploit.framework.events.on_exploit_success(exploit, session)
		end

		return session
	end

	attr_accessor :exploit # :nodoc:
	attr_accessor :payload # :nodoc:
	attr_accessor :use_job # :nodoc:
	attr_accessor :force_wait_for_session # :nodoc:

protected

	attr_accessor :session # :nodoc:

	#
	# Job run proc, sets up the exploit and kicks it off.
	#
	def job_run_proc(ctx)
		begin
			exploit, payload = ctx
		
			# Set the exploit up the bomb
			exploit.setup

			# Launch the exploit
			exploit.exploit

			# Wait the payload to acquire a session if this isn't a passive-style
			# exploit.
			if (exploit.passive? == false or force_wait_for_session == true)
				self.session = payload.wait_for_session(
					(exploit.passive? == true) ? nil : payload.wfs_delay)
			end
		rescue
			elog("Exploit failed: #{$!}", 'core', LEV_0)

			payload.stop_handler
			exploit.cleanup

			raise $!
		end
	end

	#
	# Clean up the exploit and the handler after the job completes.
	#
	def job_cleanup_proc(ctx)
		exploit, payload = ctx
		
		# 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

end

end