module Rex
module Poly

require 'rex/poly/register'
require 'rex/poly/block'

###
#
# This class encapsulates the state of a single polymorphic block set
# generation.  It tracks the current set of consumed registers, the linear
# list of blocks generated, the end-result buffer, and the phase of
# generation.  The fields exposed by the State class are intended for use only
# by the polymorphic generation subsystem and should not be modified directly.
#
###
class State

	#
	# Initializes the polymorphic generation state.
	#
	def initialize
		@block_list = nil
		reset
	end

	#
	# Resets the generation state to have a plain start by clearing all
	# consumed registers, resetting the polymorphic buffer back to its
	# beginning and destroying any block generation state.
	#
	def reset
		# Reset the generation flag on any blocks in the block list
		@block_list.each { |block|
			block[0].generated = false
		} if (@block_list)

		@regnums     = Hash.new
		@buffer      = ''
		@block_list  = []
		@curr_offset = 0
		@first_phase = true
		@badchars    = nil
	end

	#
	# Returns true if the supplied register number is already consumed.
	#
	def consumed_regnum?(regnum)
		@regnums[regnum]
	end

	#
	# Consumes a register number, thus removing it from the pool that can be
	# assigned.  The consumed register number is returned to the caller.
	#
	def consume_regnum(regnum)
		raise RuntimeError, "Register #{regnum} is already consumed." if (consumed_regnum?(regnum))

		@regnums[regnum] = true	

		regnum
	end

	#
	# Acquires a register number that has not already been consumed from the
	# supplied register number set and consumes it, returning the selected
	# register number to the caller.  The register number is selected from the
	# set at random.
	#
	def consume_regnum_from_set(regnum_set)
		# Pick a random starting point within the supplied set.
		idx = rand(regnum_set.length)

		# Try each index in the set.
		regnum_set.length.times { |x|
			regnum = regnum_set[(idx + x) % regnum_set.length]

			next if (consumed_regnum?(regnum))

			return consume_regnum(regnum)
		}

		# If we get through the entire iteration without finding a register,
		# then we are out of registers to assign.
		raise RuntimeError, "No registers are available to consume from the set"
	end

	#
	# Eliminates a register number from the consumed pool so that it can be
	# used in the future.  This happens after a block indicates that a register
	# has been clobbered.
	#
	def defecate_regnum(regnum)
		@regnums.delete(regnum)
	end

	#
	# The buffer state for the current polymorphic generation.  This stores the
	# end-result of a call to generate on a LogicalBlock.
	#
	attr_accessor :buffer

	#
	# The linear list of blocks that is generated by calling the generate
	# method on a LogicalBlock.
	#
	attr_accessor :block_list

	#
	# The current offset into the polymorphic buffer that is being generated.
	# This is updated as blocks are appended to the block_list.
	#
	attr_accessor :curr_offset

	#
	# A boolean field that is used by the LogicalBlock class to track whether
	# or not it is in the first phase (generating the block list), or in the
	# second phase (generating the polymorphic buffer).  This phases are used
	# to indicate whether or not the offset_of and regnum_of methods will
	# return actual results.
	#
	attr_accessor :first_phase

	#
	# Characters to avoid when selecting permutations, if any.
	#
	attr_accessor :badchars

end

end
end