librex 0.0.1

data/lib/rex/poly/block.rb

@@ -0,0 +1,468 @@
+module Rex
+module Poly
+
+###
+#
+# This class encapsulates a LogicalBlock permutation.  Block permutations can
+# take the form of a static string or a procedure.  This makes it possible to
+# have simple blocks and more complicated ones that take into account other
+# variables, such as dynamic registers.  The to_s method will return the
+# string version of the permutation, regardless of whether or not the
+# underlying permutation is a string or a procedure.
+#
+###
+class Permutation
+
+	#
+	# Initializes the permutation and its associated block.
+	#
+	def initialize(perm, block)
+		@perm  = perm
+		@block = block
+	end
+
+	#
+	# Returns the length of the string returned by to_s.
+	#
+	def length
+		to_s.length
+	end
+
+	#
+	# Returns the string representation of the permutation.  If the underlying
+	# permutation is a procedure, the procedure is called.  Otherwise, the
+	# string representation of the permutation is returned.
+	#
+	def to_s
+		if (@perm.kind_of?(Proc))
+			@perm.call(@block).to_s
+		else
+			@perm.to_s
+		end
+	end
+
+	attr_reader :perm
+
+end
+
+###
+#
+# This class represents a logical block which is defined as a concise portion
+# of code that may have one or more functionally equivalent implementations.
+# A logical block should serve a very specific purpose, and any permutations
+# beyond the first should result in exactly the same functionality without any
+# adverse side effects to other blocks.  
+#
+# Like blocks of code, LogicalBlock's can depend on one another in terms of
+# ordering and precedence.  By marking blocks as dependent on another, a
+# hierarchy begins to form.  This is a block dependency graph.
+#
+# To add permutations to a LogicalBlock, they can either be passed in as a
+# list of arguments to the constructor following the blocks name or can be
+# added on the fly by calling the add_perm method.  To get a random
+# permutation, the rand_perm method can be called.
+#
+# To mark one block as depending on another, the depends_on method can be
+# called with zero or more LogicalBlock instances as parameters.
+#
+###
+class LogicalBlock
+
+	#
+	# Initializes the logical block's name along with zero or more specific
+	# blocks.
+	#
+	def initialize(name, *perms)
+		@name  = name
+
+		reset
+
+		add_perm(*perms)
+	end
+
+	#
+	# Resets the block back to its starting point.
+	#
+	def reset
+		@perms           = []
+		@depends         = []
+		@next_blocks     = []
+		@clobbers        = []
+		@offset          = nil
+		@state           = nil
+		@once            = false
+		@references      = 0
+		@used_references = 0
+		@generated       = false
+	end
+
+	#
+	# Returns the block's name.
+	#
+	def name
+		@name
+	end
+
+	#
+	# Flags whether or not the block should only be generated once.  This can
+	# be used to mark a blog as being depended upon by multiple blocks, but
+	# making it such that it is only generated once.
+	#
+	def once=(tf)
+		@once = tf
+	end
+
+	#
+	# Returns true if this block is a 'once' block.  That is, this block is
+	# dependend upon by multiple blocks but should only be generated once.
+	#
+	def once
+		@once
+	end
+
+	#
+	# Increments the number of blocks that depend on this block.
+	#
+	def ref
+		@references += 1
+	end
+
+	#
+	# Increments the number of blocks that have completed their dependency
+	# pass on this block.  This number should never become higher than the
+	# @references attribute.
+	#
+	def deref
+		@used_references += 1
+	end
+
+	#
+	# Returns true if there is only one block reference remaining.
+	#
+	def last_reference?
+		(@references - @used_references <= 0)
+	end
+
+	#
+	# Adds zero or more specific permutations that may be represented either as
+	# strings or as Proc's to be called at evaluation time.
+	#
+	def add_perm(*perms)
+		@perms.concat(perms)
+	end
+
+	#
+	# Returns a random permutation that is encapsulated in a Permutation class
+	# instance.
+	#
+	def rand_perm
+		perm = nil
+
+		if (@state.badchars)
+			perm = rand_perm_badchars
+		else
+			perm = Permutation.new(@perms[rand(@perms.length)], self)
+		end
+
+		if (perm.nil?)
+			raise RuntimeError, "Failed to locate a valid permutation."
+		end
+
+		perm
+	end
+
+	#
+	# Returns a random permutation that passes any necessary bad character
+	# checks.
+	#
+	def rand_perm_badchars
+		idx = rand(@perms.length)
+		off = 0
+
+		while (off < @perms.length)
+			p = @perms[(idx + off) % @perms.length]
+
+			if (p.kind_of?(Proc) or
+			    @state.badchars.nil? or 
+			    Rex::Text.badchar_index(p, @state.badchars).nil?)
+				return Permutation.new(p, self)
+			end
+
+			off += 1
+		end
+	end
+
+	#
+	# Sets the blocks that this block instance depends on.
+	#
+	def depends_on(*depends)
+		@depends = depends.dup
+
+		# Increment dependent references
+		@depends.each { |b| b.ref }
+	end
+
+	#
+	# Defines the next blocks, but not in a dependency fashion but rather in a
+	# linking of separate block contexts.
+	#
+	def next_blocks(*blocks)
+		@next_blocks = blocks.dup
+	end
+
+	#
+	# Defines the list of zero or more LogicalRegister's that this block
+	# clobbers.
+	#
+	def clobbers(*registers)
+		@clobbers = registers
+	end
+
+	#
+	# Enumerates each register instance that is clobbered by this block.
+	#
+	def each_clobbers(&block)
+		@clobbers.each(&block)
+	end
+
+	#
+	# Generates the polymorphic buffer that results from this block and any of
+	# the blocks that it either directly or indirectly depends on.  A list of
+	# register numbers to be saved can be passed in as an argument.
+	#
+	# This method is not thread safe.  To call this method on a single block
+	# instance from within multiple threads, be sure to encapsulate the calls
+	# inside a locked context.
+	#
+	def generate(save_registers = nil, state = nil, badchars = nil)
+		# Create a localized state instance if one was not supplied.
+		state = Rex::Poly::State.new if (state == nil)
+		buf   = nil
+		cnt   = 0
+
+		# This is a lame way of doing this.  We just try to generate at most 128
+		# times until we don't have badchars.  The reason we have to do it this
+		# way is because of the fact that badchars can be introduced through
+		# block offsetting and register number selection which can't be readily
+		# predicted or detected during the generation phase.  In the future we
+		# can make this better, but for now this will have to do.
+		begin
+			buf = do_generate(save_registers, state, badchars)
+
+			if (buf and 
+			    (badchars.nil? or Rex::Text.badchar_index(buf, badchars).nil?))
+				break
+			end
+		end while ((cnt += 1) < 128)
+
+		# If we passed 128 tries, then we can't succeed.
+		buf = nil if (cnt >= 128)
+
+		buf
+	end
+	
+	#
+	# Returns the offset of a block.  If the active state for this instance is
+	# operating in the first phase, then zero is always returned.  Otherwise,
+	# the correct offset for the supplied block is returned.
+	#
+	def offset_of(lblock)
+		if (@state.first_phase) 
+			0 
+		else
+			if (lblock.kind_of?(SymbolicBlock::End))
+				@state.curr_offset
+			else
+				lblock.offset
+			end
+		end
+	end
+
+	#
+	# Returns the register number associated with the supplied LogicalRegister
+	# instance.  If the active state for this instance is operating in the
+	# first phase, then zero is always returned.  Otherwise, the correct
+	# register number is returned based on what is currently assigned to the
+	# supplied LogicalRegister instance, if anything.
+	#
+	def regnum_of(reg)
+		(@state.first_phase) ? 0 : reg.regnum
+	end
+
+	#
+	# This attributes contains the currently assigned offset of the permutation
+	# associated with this block into the polymorphic buffer that is being
+	# generated.
+	#
+	attr_accessor :offset
+
+	#
+	# Whether or not this block has currently been generated for a given
+	# iteration.
+	#
+	attr_accessor :generated
+
+protected
+	
+	#
+	# Performs the actual polymorphic buffer generation.  Called from generate
+	#
+	def do_generate(save_registers, state, badchars)
+		# Reset the state in case it was passed in.
+		state.reset
+
+		# Set the bad character list
+		state.badchars = badchars if (badchars)
+
+		# Consume any registers that should be saved.
+		save_registers.each { |reg|
+			state.consume_regnum(reg)
+		} if (save_registers)
+
+		# Build the linear list of blocks that will be processed.  This
+		# list is built in a dynamic fashion based on block dependencies.
+		# The list that is returned is an Array of which each element is a two
+		# member array, the first element being the LogicalBlock instance that
+		# the permutation came from and the second being an instance of the
+		# Permutation class associated with the selected permutation.
+		block_list = generate_block_list(state)
+
+		# Transition into the second phase which enables offset_of and regnum_of
+		# calls to return real values.
+		state.first_phase = false
+
+		# Now that every block has been assigned an offset, generate the
+		# buffer block by block, assigning registers as necessary.
+		block_list.each { |b| 
+
+			# Generate the next permutation and append it to the buffer.
+			begin
+				state.buffer += b[1].to_s 
+			# If an invalid register exception is raised, try to consume a random
+			# register from the register's associated architecture register
+			# number set.
+			rescue InvalidRegisterError => e
+				e.reg.regnum = state.consume_regnum_from_set(e.reg.class.regnum_set)
+				retry
+			end
+	
+			# Remove any of the registers that have been clobbered by this block
+			# from the list of consumed register numbers so that they can be used
+			# in the future.
+			b[0].each_clobbers { |reg|
+				begin
+					state.defecate_regnum(reg.regnum)
+
+					reg.regnum = nil
+				rescue InvalidRegisterError
+				end
+			}
+				
+		}
+
+		# Finally, return the buffer that has been created.
+		state.buffer
+	end
+
+	#
+	# Generates the linear list of block permutations which is stored in the
+	# supplied state instance.  This is done prior to assigning blocks offsets
+	#
+	def generate_block_list(state, level=0)
+		if @depends.length > 1
+			@depends.length.times {
+				f = rand(@depends.length)
+				@depends.push(@depends.delete_at(f))
+			}
+		end
+
+		@depends.length.times { |cidx|
+		
+			pass = false
+		
+			while (not pass)
+			
+				if (@depends[cidx].generated)
+					break
+
+				# If this dependent block is a once block and the magic 8 ball turns
+				# up zero, skip it and let a later block pick it up.  We only do this
+				# if we are not the last block to have a dependency on this block.
+				elsif ((@depends[cidx].once) and
+			    	(rand(2).to_i == 0) and
+			    	(@depends[cidx].last_reference? == false))
+					break
+				end
+
+				# Generate this block
+				@depends[cidx].generate_block_list(state, level+1)
+
+				if level != 0
+					return
+				else
+					@depends.length.times {
+						f = rand(@depends.length)
+						@depends.push(@depends.delete_at(f))
+					}
+
+					next
+				end
+			end
+				
+			next
+		}
+
+		self.deref
+
+		# Assign the instance local state for the duration of this generation
+		@state = state
+
+		# Select a random permutation
+		perm = rand_perm
+
+		# Set our block offset to the current state offset
+		self.offset = state.curr_offset
+
+		# Flag ourselves as having been generated for this iteration.
+		self.generated = true
+
+		# Adjust the current offset based on the permutations length
+		state.curr_offset += perm.length
+
+		# Add it to the linear list of blocks
+		state.block_list << [ self, perm ]
+
+		# Generate all the blocks that follow this one.
+		@next_blocks.each { |b|
+			b.generate_block_list(state)
+		}
+
+		# Return the state's block list
+		state.block_list
+	end
+
+end
+
+###
+#
+# Symbolic blocks are used as special-case LogicalBlock's that have meaning
+# a more general meaning.  For instance, SymbolicBlock::End can be used to
+# symbolize the end of a polymorphic buffer.
+#
+###
+module SymbolicBlock
+
+	###
+	#
+	# The symbolic end of a polymorphic buffer.
+	#
+	###
+	class End < LogicalBlock
+		def initialize
+			super('__SYMBLK_END__')
+		end
+	end
+end
+
+end
+end