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.job_id                 = nil
		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)
		# Try to use the target's platform in preference of the exploit's
		exp_platform = exploit.targets[target_idx].platform || exploit.platform

		return ((payload.platform & exp_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

		# If we are being instructed to run as a job then let's create that job
		# like a good person.
		if (use_job or exploit.stance == Msf::Exploit::Stance::Passive)
			# Since references to the exploit and payload will hang around for
			# awhile in the job, make sure we copy them so further changes to
			# the datastore and so forth don't alter settings in existing jobs
			#
			# This doesn't actually work.  For some reason the payload and the
			# exploit don't match up properly after the copy.  Disable for now
			# and revert to the previous behavior.
			e = exploit.replicant
			p = payload.replicant
			ctx = [ e, p ]
			#ctx = [ exploit, payload ]

			self.job_id  = exploit.framework.jobs.start_bg_job(
				"Exploit: #{exploit.refname}",
				ctx,
				Proc.new { |ctx_| job_run_proc(ctx_) },
				Proc.new { |ctx_| job_cleanup_proc(ctx_) }
			)
		else
			# No need to copy since we aren't creating a job.  We wait until
			# they're finished running to do anything else with them, so
			# nothing should be able to modify their datastore or other
			# settings until after they're done.
			ctx = [ exploit, payload ]
			job_run_proc(ctx)
			job_cleanup_proc(ctx)
		end

		return session
	end

	attr_accessor :exploit # :nodoc:
	attr_accessor :payload # :nodoc:
	attr_accessor :use_job # :nodoc:
	#
	# The identifier of the job this exploit is launched as, if it's run as a
	# job.
	#
	attr_accessor :job_id
	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

			exploit.framework.events.on_module_run(exploit)

			# 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 + exploit.wfs_delay)
			end
		rescue ::Exception
			exploit.framework.events.on_module_error(exploit, $!)
			exploit.print_error("Exploit failed: #{$!}")
			exploit.error = $!

			elog("Exploit failed (#{exploit.refname}): #{$!.class} #{$!}", 'core', LEV_0)
			dlog("Call stack:\n#{$@.join("\n")}", 'core', LEV_3)

			payload.stop_handler
			exploit.cleanup
			return
		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

		exploit.framework.events.on_module_complete(exploit)

		# Allow the exploit to cleanup after itself, that messy bugger.
		exploit.cleanup
	end

end

end