#!/usr/bin/env ruby

module Rex
module Arch

#
# everything here is mostly stole from vlad's perl x86 stuff
#

module X86

	#
	# Register number constants
	#
	EAX = AL = AX = ES = 0
	ECX = CL = CX = CS = 1
	EDX = DL = DX = SS = 2
	EBX = BL = BX = DS = 3
	ESP = AH = SP = FS = 4
	EBP = CH = BP = GS = 5
	ESI = DH = SI =      6
	EDI = BH = DI =      7

	REG_NAMES32 = [ 'eax', 'ecx', 'edx', 'ebx', 
	                'esp', 'ebp', 'esi', 'edi' ] # :nodoc:

	# Jump tp a specific register
	def self.jmp_reg(str)
		reg = reg_number(str)
		_check_reg(reg)
		"\xFF" + [224 + reg].pack('C')
	end

	# This method returns the opcodes that compose a jump instruction to the
	# supplied relative offset.
	def self.jmp(addr)
		"\xe9" + pack_dword(rel_number(addr))
	end

	#
	# This method adds/subs a packed long integer
	#
	def self.dword_adjust(dword, amount=0)
		[dword.unpack('V')[0] + amount].pack('V')
	end
	
	#
	# This method returns the opcodes that compose a tag-based search routine
	#
	def self.searcher(tag)
		"\xbe" + dword_adjust(tag,-1)+  # mov esi, Tag - 1
		"\x46" +                        # inc esi
		"\x47" +                        # inc edi (end_search:)
		"\x39\x37" +                    # cmp [edi],esi
		"\x75\xfb" +                    # jnz 0xa (end_search)
		"\x46" +                        # inc esi
		"\x4f" +                        # dec edi (start_search:)
		"\x39\x77\xfc" +                # cmp [edi-0x4],esi
		"\x75\xfa" +                    # jnz 0x10 (start_search)
		jmp_reg('edi')                  # jmp edi	
	end

	#
	# Generates a buffer that will copy memory immediately following the stub
	# that is generated to be copied to the stack
	#
	def self.copy_to_stack(len)
		# four byte align
		len = (len + 3) & ~0x3

		stub =
			"\xeb\x0f"+                # jmp _end
			"\x68" + [len].pack('V')+  # push n
			"\x59"+                    # pop ecx
			"\x5e"+                    # pop esi
			"\x29\xcc"+                # sub esp, ecx
			"\x89\xe7"+                # mov edi, esp
			"\xf3\xa4"+                # rep movsb
			"\xff\xe4"+                # jmp esp
			"\xe8\xec\xff\xff\xff"     # call _start

		stub
	end
	
	#
	# This method returns the opcodes that compose a short jump instruction to
	# the supplied relative offset.
	#
	def self.jmp_short(addr)
		"\xeb" + pack_lsb(rel_number(addr, -2))
	end

	#
	# This method returns the opcodes that compose a relative call instruction
	# to the address specified.
	#
	def self.call(addr)
		"\xe8" + pack_dword(rel_number(addr, -5))
	end

	#
	# This method returns a number offset to the supplied string.
	#
	def self.rel_number(num, delta = 0)
		s = num.to_s

		case s[0, 2]
			when '$+'
				num = s[2 .. -1].to_i
			when '$-'
				num = -1 * s[2 .. -1].to_i
			when '0x'
				num = s.hex
			else
				delta = 0
		end

		return num + delta
	end

	#
	# This method returns the number associated with a named register.
	#
	def self.reg_number(str)
		return self.const_get(str.upcase)
	end

	#
	# This method returns the register named associated with a given register
	# number.
	#
	def self.reg_name32(num)
		_check_reg(num)
		return REG_NAMES32[num].dup
	end

	#
	# This method generates the encoded effective value for a register.
	#
	def self.encode_effective(shift, dst)
		return (0xc0 | (shift << 3) | dst)
	end

	#
	# This method generates the mod r/m character for a source and destination
	# register.
	#
	def self.encode_modrm(dst, src)
		_check_reg(dst, src)
		return (0xc0 | src | dst << 3).chr
	end

	#
	# This method generates a push byte instruction.
	#
	def self.push_byte(byte)
		# push byte will sign extend...
		if byte < 128 && byte >= -128
			return "\x6a" + (byte & 0xff).chr
		end
		raise ::ArgumentError, "Can only take signed byte values!", caller()
	end

	#
	# This method generates a push word instruction.
	#
	def self.push_word(val)
		return "\x66\x68" + [ val ].pack('v') 
	end

	#
	# This method generates a push dword instruction.
	#
	def self.push_dword(val)
		return "\x68" + [ val ].pack('V') 
	end

	#
	# This method generates a pop dword instruction into a register.
	#
	def self.pop_dword(dst)
		_check_reg(dst)
		return (0x58 | dst).chr
	end

	#
	# This method generates an instruction that clears the supplied register in
	# a manner that attempts to avoid bad characters, if supplied.
	#
	def self.clear(reg, badchars = '')
		_check_reg(reg)
		return set(reg, 0, badchars)
	end

	#
	# This method generates the opcodes that set the low byte of a given
	# register to the supplied value.
	#
	def self.mov_byte(reg, val)
		_check_reg(reg)
		# chr will raise RangeError if val not between 0 .. 255
		return (0xb0 | reg).chr + val.chr
	end

	#
	# This method generates the opcodes that set the low word of a given
	# register to the supplied value.
	#
	def self.mov_word(reg, val)
		_check_reg(reg)
		if val < 0 || val > 0xffff
			raise RangeError, "Can only take unsigned word values!", caller()
		end
		return "\x66" + (0xb8 | reg).chr + [ val ].pack('v')
	end

	#
	# This method generates the opcodes that set the a register to the
	# supplied value.
	#
	def self.mov_dword(reg, val)
		_check_reg(reg)
		return (0xb8 | reg).chr + [ val ].pack('V')
	end

	#
	# This method is a general way of setting a register to a value.  Depending
	# on the value supplied, different sets of instructions may be used.
	#
	# TODO: Make this moderatly intelligent so it chain instructions by itself
		#   (ie. xor eax, eax + mov al, 4 + xchg ah, al)
	def self.set(dst, val, badchars = '')
		_check_reg(dst)

		# If the value is 0 try xor/sub dst, dst (2 bytes)
		if(val == 0)
			opcodes = Rex::Text.remove_badchars("\x29\x2b\x31\x33", badchars)
			if !opcodes.empty?
				return opcodes[rand(opcodes.length)].chr + encode_modrm(dst, dst)
			end
# TODO: SHL/SHR
# TODO: AND 
		end

		# try push BYTE val; pop dst (3 bytes)
		begin
			return _check_badchars(push_byte(val) + pop_dword(dst), badchars)
		rescue ::ArgumentError, ::RuntimeError, ::RangeError
		end

		# try clear dst, mov BYTE dst (4 bytes)
		begin
			# break if val == 0
			return _check_badchars(clear(dst, badchars) + mov_byte(dst, val), badchars)
		rescue ::ArgumentError, ::RuntimeError, ::RangeError
		end

		# try mov DWORD dst (5 bytes)
		begin
			return _check_badchars(mov_dword(dst, val), badchars)
		rescue ::ArgumentError, ::RuntimeError, ::RangeError
		end

		# try push DWORD, pop dst (6 bytes)
		begin
			return _check_badchars(push_dword(val) + pop_dword(dst), badchars)
		rescue ::ArgumentError, ::RuntimeError, ::RangeError
		end

		# try clear dst, mov WORD dst (6 bytes)
		begin
			# break if val == 0
			return _check_badchars(clear(dst, badchars) + mov_word(dst, val), badchars)
		rescue ::ArgumentError, ::RuntimeError, ::RangeError
		end

		raise RuntimeError, "No valid set instruction could be created!", caller()
	end

	#
	# Builds a subtraction instruction using the supplied operand
	# and register.
	#
	def self.sub(val, reg, badchars = '', add = false, adjust = false, bits = 0)
		opcodes = []
		shift   = (add == true) ? 0 : 5

		if (bits <= 8 and val >= -0x7f and val <= 0x7f)
			opcodes << 
				((adjust) ? '' : clear(reg, badchars)) + 
				"\x83" + 
				[ encode_effective(shift, reg) ].pack('C') +
				[ val.to_i ].pack('C')
		end

		if (bits <= 16 and val >= -0xffff and val <= 0)
			opcodes << 
				((adjust) ? '' : clear(reg, badchars)) + 
				"\x66\x81" + 
				[ encode_effective(shift, reg) ].pack('C') +
				[ val.to_i ].pack('v')
		end
			
		opcodes << 
			((adjust) ? '' : clear(reg, badchars)) + 
			"\x81" + 
			[ encode_effective(shift, reg) ].pack('C') +
			[ val.to_i ].pack('V')

		# Search for a compatible opcode
		opcodes.each { |op|
			begin 
				_check_badchars(op, badchars)
			rescue
				next
			end

			return op
		}

		if opcodes.empty?
			raise RuntimeError, "Could not find a usable opcode", caller()
		end
	end

	#
	# This method generates the opcodes equivalent to subtracting with a
	# negative value from a given register.
	#
	def self.add(val, reg, badchars = '', adjust = false, bits = 0)
		sub(val, reg, badchars, true, adjust, bits)
	end

	#
	# This method wrappers packing an integer as a little-endian buffer.
	#
	def self.pack_dword(num)
		[num].pack('V')
	end

	#
	# This method returns the least significant byte of a packed dword.
	#
	def self.pack_lsb(num)
		pack_dword(num)[0,1]
	end

	#
	# This method adjusts the value of the ESP register by a given amount.
	#
	def self.adjust_reg(reg, adjustment)
		if (adjustment > 0)
			sub(adjustment, reg, '', false, false, 32)
		else
			add(adjustment, reg, '', true, 32)
		end
	end

	def self._check_reg(*regs) # :nodoc:
		regs.each { |reg|
			if reg > 7 || reg < 0
				raise ArgumentError, "Invalid register #{reg}", caller()
			end
		}
		return nil
	end

	def self._check_badchars(data, badchars) # :nodoc:
		idx = Rex::Text.badchar_index(data, badchars)
		if idx
			raise RuntimeError, "Bad character at #{idx}", caller()
		end
		return data
	end

	#
	# This method returns an array of 'safe' FPU instructions
	#
	def self.fpu_instructions
		fpus = []

		0xe8.upto(0xee) { |x| fpus << "\xd9" + x.chr }
		0xc0.upto(0xcf) { |x| fpus << "\xd9" + x.chr }
		0xc0.upto(0xdf) { |x| fpus << "\xda" + x.chr }
		0xc0.upto(0xdf) { |x| fpus << "\xdb" + x.chr }
		0xc0.upto(0xc7) { |x| fpus << "\xdd" + x.chr }

		fpus << "\xd9\xd0"
		fpus << "\xd9\xe1"
		fpus << "\xd9\xf6"
		fpus << "\xd9\xf7"
		fpus << "\xd9\xe5"

		# This FPU instruction seems to fail consistently on Linux
		#fpus << "\xdb\xe1"

		fpus
	end

	#
	# This method returns an array containing a geteip stub, a register, and an offset
	# This method will return nil if the getip generation fails
	#
	def self.geteip_fpu(badchars)
		
		#
		# Default badchars to an empty string
		#
		badchars ||= ''
		
		#
		# Bail out early if D9 is restricted
		#
		return nil if badchars.index("\xd9")

		#
		# Create a list of FPU instructions
		#
		fpus = *self.fpu_instructions
		bads = []
		badchars.each_byte  do |c|
			fpus.each do |str|
				bads << str if (str.index(c.chr))
			end
		end	
		bads.each { |str| fpus.delete(str) }
		return nil if fpus.length == 0

		#
		# Create a list of registers to use for fnstenv
		#
		dsts = []
		0.upto(7) do |c|
			dsts << c if (not badchars.index( (0x70+c).chr ))
		end

		if (dsts.include?(ESP) and badchars.index("\x24"))
			dsts.delete(ESP)
		end

		return nil if dsts.length == 0

		#
		# Grab a random FPU instruction
		#
		fpu = fpus[ rand(fpus.length) ]

		#
		# Grab a random register from dst
		#
		while(dsts.length > 0)
			buf = ''
			dst = dsts[ rand(dsts.length) ]
			dsts.delete(dst)

			# If the register is not ESP, copy ESP
			if (dst != ESP)
				next if badchars.index( (0x70 + dst).chr )

				if (not (badchars.index("\x89") or badchars.index( (0xE0+dst).chr )))
					buf << "\x89" + (0xE0 + dst).chr
				else
					next if badchars.index("\x54")
					next if badchars.index( (0x58+dst).chr )
					buf << "\x54" + (0x58 + dst).chr
				end
			end

			pad = 0
			while (pad < (128-12) and badchars.index( (256-12-pad)))
				pad += 4
			end

			# Give up on finding a value to use here
			if (pad == (128-12))
				return nil
			end

			out = buf + fpu + "\xd9" + (0x70 + dst).chr
			out << "\x24" if dst == ESP
			out << (256-12-pad).chr

			regs = [*(0..7)]
			while (regs.length > 0)
				reg = regs[ rand(regs.length) ]
				regs.delete(reg)
				next if reg == ESP
				next if badchars.index( (0x58 + reg).chr )

				# Pop the value back out
				0.upto(pad / 4) { |c| out << (0x58 + reg).chr }

				# Fix the value to point to self
				gap = out.length - buf.length

				return [out, REG_NAMES32[reg].upcase, gap]
			end
		end

		return nil
	end

end

end end