require 'rex/arch/x86'
require 'rex/nop/opty2_tables'

module Rex
module Nop

###
#
# This class provides an interface to generating multi-byte NOP sleds for x86.
# Optyx and spoonm get the creds!
#
###
class Opty2

	Table = Rex::Nop::Opty2Tables::StateTable

	def initialize(badchars = '', save_registers = nil)
		self.badchars       = badchars
		self.save_registers = (save_registers || []) | [ 'esp', 'ebp']
	end

	#
	# Generates the Opty2 multi-byte NOP sled.
	#
	def generate_sled(length)
		return '' if (length <= 0)

		# Initialize the sled buffer, the previous state, and the current stream
		# length.
		sled = ''
		prev = 256
		slen = 0

		# Initialize the byte count array
		counts = []

		256.times { |idx| counts[idx] = 0 }

		# Initialize the bad register mask
		mask = 0

		save_registers.each { |reg|
			mask |= 1 << (Rex::Arch::X86.reg_number(reg))
		}
		mask = mask << 16

		# Initialize the bad byte lookup table
		bad_bytes = []
		(badchars || '').each_byte { |byte|
			bad_bytes[byte] = 1
		}

		# Build the sled
		while (length > 0)
			low  = -1
			lows = []

			Table[prev].each { |nt|
				nt.each { |e|
					# Skip it if it's masked off or too large
					next if ((e & mask) != 0)
					next if (((e >> 8) & 0xff) > slen)
					
					byte = e & 0xff

					# Skip it if it's a bad byte
					next if (bad_bytes[byte] == 1)

					# Use it if it's a better value
					if ((low == -1) or (low > counts[byte]))
						low  = counts[byte]
						lows = [byte]
					# Otherwise, if it's just as good..
					elsif (low == counts[byte])
						lows << byte
					end
				}
			}

			# If we didn't find at least one byte possibility, then we're stuck.
			# Abort.
			if (low == -1)
				raise RuntimeError, "Failed to find a valid byte."
			end

			# Pick a random character for the possiblities
			prev = lows[rand(lows.length)]

			# Increment its used count
			counts[prev] += 1

			# Prepend the byte to the sled
			sled = prev.chr + sled

			# Increment the sled length
			slen   += 1
			length -= 1
		end

		# Return the sled
		sled
	end

	attr_accessor :badchars, :save_registers # :nodoc:
end

end
end