##
# $Id$
##

###
#
# This module exposes a simple method to create an payload in an executable.
#
###

module Msf
module Exploit::EXE

	def initialize(info = {})
		super

		register_advanced_options(
			[
				OptString.new( 'EXE::Custom',   [ false, 'Use custom exe instead of automatically generating a payload exe']),
				OptString.new( 'EXE::Path',     [ false, 'The directory in which to look for the executable template' ]),
				OptString.new( 'EXE::Template', [ false, 'The executable template file name.' ]),
				OptBool.new(   'EXE::Inject',   [ false, 'Set to preserve the original EXE function' ]),
				OptBool.new(   'EXE::OldMethod', [ false, 'Set to use the substitution EXE generation method.' ]),
				OptBool.new(   'EXE::FallBack', [ false, 'Use the default template in case the specified one is missing' ])
			], self.class)
	end

	def get_custom_exe
		datastore['DisablePayloadHandler'] = true
		file = ::File.open(datastore['EXE::Custom'],'rb')
		exe = file.read(file.stat.size)
		file.close
		exe
	end

	def generate_payload_exe(opts = {})
		return get_custom_exe if datastore.include? 'EXE::Custom'

		exe_init_options(opts)

		pl = opts[:code]
		pl ||= payload.encoded

		# Fall back to x86...
		if not opts[:arch] or opts[:arch].length < 1
			opts[:arch] = [ ARCH_X86 ]
		end
		# Ensure we have an array
		if not opts[:arch].kind_of? Array
			opts[:arch] = [ opts[:arch] ]
		end

		# Transform the PlatformList
		if (opts[:platform].kind_of? Msf::Module::PlatformList)
			opts[:platform] = opts[:platform].platforms
		end

		exe = Msf::Util::EXE.to_executable(framework, opts[:arch], opts[:platform], pl, opts)
		exe_post_generation(opts)
		exe
	end

	def generate_payload_exe_service(opts = {})
		return get_custom_exe if datastore.include? 'EXE::Custom'

		exe_init_options(opts)

		# NOTE: Only Windows is supported here.
		pl = opts[:code]
		pl ||= payload.encoded

		if opts[:arch] and opts[:arch] == ARCH_X64
			exe = Msf::Util::EXE.to_win64pe_service(framework, pl, opts)
		else
			exe = Msf::Util::EXE.to_win32pe_service(framework, pl, opts)
		end

		exe_post_generation(opts)
		exe
	end

	def generate_payload_dll(opts = {})
		return get_custom_exe if datastore.include? 'EXE::Custom'

		exe_init_options(opts)

		# NOTE: Only Windows is supported here.
		pl = opts[:code]
		pl ||= payload.encoded

		if opts[:arch] and opts[:arch] == ARCH_X64
			dll = Msf::Util::EXE.to_win64pe_dll(framework, pl, opts)
		else
			dll = Msf::Util::EXE.to_win32pe_dll(framework, pl, opts)
		end

		exe_post_generation(opts)
		dll
	end

protected
	def exe_init_options(opts)
		opts.merge!(
			{
				:template_path => datastore['EXE::Path'],
				:template => datastore['EXE::Template'],
				:inject => datastore['EXE::Inject'],
				:fallback => datastore['EXE::FallBack'],
				:sub_method => datastore['EXE::OldMethod']
			})

		# Prefer the target's platform/architecture information, but use
		# the module's if no target specific information exists
		opts[:platform] ||= target_platform  if self.respond_to? :target_platform
		opts[:platform] ||= platform         if self.respond_to? :platform
		opts[:arch] ||= target_arch          if self.respond_to? :target_arch
		opts[:arch] ||= arch                 if self.respond_to? :arch
	end

	def exe_post_generation(opts)
		if (opts[:fellback])
			print_status("Warning: Falling back to default template: #{opts[:fellback]}")
		end
	end

end
end