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

		# Default the session to nil
		self.session = nil

		# Explicitly clear the module's job_id in case it was set in a previous
		# run
		exploit.job_id = 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 don't alter settings in existing jobs
			e = exploit.replicant
			p = payload.replicant

			# Assign the correct exploit instance to the payload
			p.assoc_exploit = e

			# Generate the encoded version of the supplied payload for the
			# newly copied exploit module instance
			e.generate_payload(p)
			ctx = [ e, p ]

			e.job_id = e.framework.jobs.start_bg_job(
				"Exploit: #{e.refname}",
				ctx,
				Proc.new { |ctx_| job_run_proc(ctx_) },
				Proc.new { |ctx_| job_cleanup_proc(ctx_) }
			)
			self.job_id = e.job_id
		else
			# Generate the encoded version of the supplied payload on the
			# exploit module instance
			exploit.generate_payload(payload)

			# 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:
	attr_accessor :session # :nodoc:

protected



	#
	# Job run proc, sets up the exploit and kicks it off.
	#
	def job_run_proc(ctx)
		# Default session wait time..
		delay = payload.wfs_delay + exploit.wfs_delay
		delay = nil if exploit.passive?

		begin
			exploit, payload = ctx

			# Set the exploit up the bomb
			exploit.setup

			exploit.framework.events.on_module_run(exploit)

			# Launch the exploit
			exploit.exploit

		rescue ::Exception => e
			if [::RuntimeError, ::Interrupt].include?(e.class)
				# Wait for session, but don't wait long.
				delay = 0.01
			end

			# Build a user-friendly error message
			msg = "#{e}"
			msg << "#{e.class}" if msg.empty?

			# Record the error to various places
			exploit.framework.events.on_module_error(exploit, msg)
			exploit.print_error("Exploit exception: #{msg}")
			exploit.error = msg

			# Log a more verbose version
			elog("Exploit exception (#{exploit.refname}): #{e.class}: #{e}", 'core', LEV_0)
			dlog("Call stack:\n#{$@.join("\n")}", 'core', LEV_3)
		end

		# Wait the payload to acquire a session if this isn't a passive-style
		# exploit.
		return if not delay

		if (force_wait_for_session == true) or
			(exploit.passive? == false and exploit.handler_enabled?)
			begin
				self.session = payload.wait_for_session(delay)
			rescue ::Interrupt
				# Don't let interrupt pass upward
			end
		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