require 'msf/core'

module Msf

###
#
# The generic payloads are used to define a generalized payload type that
# is both architecture and platform independent.  Under the hood, generic
# payloads seek out the correct payload for the appropriate architecture
# and platform that is being targeted.
#
###
module Payload::Generic

	#
	# Registers options that are common to all generic payloads, such as
	# platform and arch.
	#
	def initialize(info = {})
		super(merge_info(info,
			'Arch'     => ARCH_ALL,
			'Platform' => ''))

		register_advanced_options(
			[
				OptString.new('PLATFORM',
					[
						false,
						"The platform that is being targeted",
						nil
					]),
				OptString.new('ARCH',
					[
						false,
						"The architecture that is being targeted",
						nil
					])
			], Msf::Payload::Generic)
	end

	#
	# Reset's the generic payload's internal state so that it can find a new
	# actual payload.
	#
	def reset
		self.explicit_arch     = nil
		self.explicit_platform = nil
		self.actual_payload    = nil
	end

	#
	# Generate is different from other methods -- it will try to re-detect
	# the actual payload in case settings have changed.  Other methods will
	# use the cached version if possible.
	#
	def generate
		reset

		redirect_to_actual(:generate)
	end

	#
	# Overrides -- we have to redirect all potential payload methods
	# to the actual payload so that they get handled appropriately, cuz
	# we're like a proxy and stuff.  We can't use method_undefined
	# because all of these methods are actually defined.
	#
	
	def payload
		redirect_to_actual(:payload)
	end

	def offsets
		redirect_to_actual(:offsets)
	end

	def substitute_vars(*args)
		redirect_to_actual(:substitute_vars, *args)
	end

	def replace_var(*args)
		redirect_to_actual(:replace_var, *args)
	end

	def compatible_encoders
		redirect_to_actual(:compatible_encoders)
	end
	
	def compatible_nops
		redirect_to_actual(:compatible_nops)
	end

	def handle_connection(*args)
		redirect_to_actual(:handle_connection, *args)
	end

	def on_session(*args)
		redirect_to_actual(:on_session, *args)
	end

	#
	# Stager overrides
	#
	
	def stage_payload
		redirect_to_actual(:stage_payload)
	end

	def stage_offsets
		redirect_to_actual(:stage_offsets)
	end

	def stager_payload
		redirect_to_actual(:stager_payload)
	end

	def stager_offsets
		redirect_to_actual(:stager_offsets)
	end

	def stage_over_connection?
		redirect_to_actual(:stage_over_connection?)
	end

	def generate_stage
		redirect_to_actual(:generate_stage)
	end

	def handle_connection_stage(*args)
		redirect_to_actual(:handle_connection_stage, *args)
	end

	def handle_intermediate_stage(*args)
		redirect_to_actual(:handle_intermediate_stage, *args)
	end

	def stage_prefix
		redirect_to_actual(:stage_prefix)
	end

	def stage_prefix=(*args)
		redirect_to_actual(:stage_prefix=, *args)
	end

	#
	# First, find the underlying payload and then pass all methods onto it.
	#
	def redirect_to_actual(name, *args)
		find_actual_payload

		actual_payload.send(name, *args)
	end

	#
	# Accessor that makes it possible to define an explicit platform.  This is
	# used for things like payload regeneration.
	#
	attr_accessor :explicit_platform
	#
	# Accessor that makes it possible to define an explicit architecture.  This
	# is used for things like payload regeneration.
	#
	attr_accessor :explicit_arch

protected

	#
	# The actual underlying platform/arch-specific payload instance that should 
	# be used.
	#
	attr_accessor :actual_payload

	#
	# Returns the actual platform that should be used for the payload.
	#
	def actual_platform
		platform = nil

		if explicit_platform.nil? == false
			platform = explicit_platform
		elsif datastore['PLATFORM']
			platform = datastore['PLATFORM']
		elsif assoc_exploit
			platform = assoc_exploit.target_platform
		end

		# If we still have an invalid platform, then we suck.
		if platform.nil?
			raise NoCompatiblePayloadError, "A platform could not be determined by the generic payload"
		elsif platform.kind_of?(String)
			platform = Msf::Module::PlatformList.transform(platform)
		end

		return platform
	end

	#
	# Returns the actual architecture that should be used for the payload.
	#
	def actual_arch
		arch = nil

		if explicit_arch.nil? == false
			arch = explicit_arch
		elsif datastore['ARCH']
			arch = datastore['ARCH']
		elsif assoc_exploit
			arch = assoc_exploit.target_arch || ARCH_X86
		end

		# If we still have an invalid architecture, then we suck.
		if arch.nil?
			raise NoCompatiblePayloadError, "An architecture could not be determined by the generic payload"
		elsif arch.kind_of?(String)
			arch = [ arch ]
		end

		return arch
	end

	def find_actual_payload
		if actual_payload.nil?
			self.actual_payload = framework.payloads.find_payload(
				actual_platform, 
				actual_arch,
				handler_klass,
				session,
				payload_type)

			if actual_payload.nil?
				raise NoCompatiblePayloadError, "Could not locate a compatible payload for #{actual_platform.names.join("/")}/#{actual_arch}"
			else
				dlog("Selected payload #{actual_payload.refname} from generic payload #{refname}", 'core', LEV_2)
				# Share our datastore with the actual payload so that it has the
				# appropriate values to substitute ad so on.
				self.actual_payload.share_datastore(self.datastore)

				# Set the associated exploit for the payload.
				self.actual_payload.assoc_exploit  = self.assoc_exploit

				# Set the parent payload to this payload so that we can handle
				# things like session creation (so that event notifications will
				# work properly)
				self.actual_payload.parent_payload = self
			end
		end
	end


end

end