# -*- coding: binary -*-
module Rex
module Poly

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

###
#
# 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