# -*- coding: binary -*-

module Rex

  module Poly

    #
    # A subclass to represent a Rex poly machine on the x86 architecture.
    #
    class MachineX86 < Rex::Poly::Machine

      def initialize( badchars='', consume_base_pointer=nil, consume_stack_pointer=true )
        super( badchars, Metasm::Ia32.new )

        @reg_available << Rex::Arch::X86::EAX
        @reg_available << Rex::Arch::X86::EBX
        @reg_available << Rex::Arch::X86::ECX
        @reg_available << Rex::Arch::X86::EDX
        @reg_available << Rex::Arch::X86::ESI
        @reg_available << Rex::Arch::X86::EDI
        @reg_available << Rex::Arch::X86::EBP
        @reg_available << Rex::Arch::X86::ESP

        # By default we consume the EBP register if badchars contains \x00. This helps speed
        # things up greatly as many instructions opperating on EBP introduce a NULL byte. For
        # example, a MOV instruction with EAX as the source operand is as follows:
        #     8B08    mov ecx, [eax]
        # but the same instruction with EBP as the source operand is as follows:
        #     8B4D00  mov ecx, [ebp] ; This is assembled as 'mov ecx, [ebp+0]'
        # we can see that EBP is encoded differently with an offset included. We can still
        # try to generate a solution with EBP included and \x00 in the badchars list but
        # it can take considerably longer.
        if( ( consume_base_pointer.nil? and not Rex::Text.badchar_index( "\x00", @badchars ).nil? ) or consume_base_pointer == true )
          create_variable( 'base_pointer', 'ebp' )
        end

        # By default we consume the ESP register to avoid munging the stack.
        if( consume_stack_pointer )
          create_variable( 'stack_pointer', 'esp' )
        end

        # discover all the safe FPU instruction we can use.
        @safe_fpu_instructions = ::Array.new
        Rex::Arch::X86.fpu_instructions.each do | fpu |
          if( is_valid?( fpu ) )
            @safe_fpu_instructions << fpu
          end
        end
      end

      #
      # The general purpose registers are 32bit
      #
      def native_size
        Rex::Poly::Machine::DWORD
      end

      #
      # Overload this method to intercept the 'set' primitive with the 'location' keyword
      # and create the block with the '_set_variable_location'. We do this to keep a
      # consistent style.
      #
      def create_block_primitive( block_name, primitive_name, *args )
        if( primitive_name == 'set' and args.length == 2 and args[1] == 'location' )
          _create_block_primitive( block_name, '_set_variable_location', args[0] )
        else
          super
        end
      end

      #
      # XXX: If we have a loop primitive, it is a decent speed bump to force the associated variable
      # of the first loop primitive to be assigned as ECX (for the x86 LOOP instruction), this is not
      # neccasary but can speed generation up significantly.
      #
      #def generate
      #	@blocks.each_value do | block |
      #		if( block.first.primitive == 'loop' )
      #			@variables.delete( block.first.args.first )
      #			create_variable( block.first.args.first, 'ecx' )
      #			break
      #		end
      #	end
      #	# ...go go go
      #	super
      #end

      protected

      #
      # Resolve a register number into a suitable register name.
      #
      def _register_value( regnum, size=nil )
        value = nil
        # we default to a native 32 bits if no size is specified.
        if( size.nil? )
          size = native_size()
        end

        if( size == Rex::Poly::Machine::DWORD )
          value = Rex::Arch::X86::REG_NAMES32[ regnum ]
        elsif( size == Rex::Poly::Machine::WORD )
          value = Rex::Arch::X86::REG_NAMES16[ regnum ]
        elsif( size == Rex::Poly::Machine::BYTE )
          # (will return nil for ESI,EDI,EBP,ESP)
          value = Rex::Arch::X86::REG_NAMES8L[ regnum ]
        else
          raise RuntimeError, "Register number '#{regnum}' (size #{size.to_i}) is unavailable."
        end
        return value
      end

      #
      # Create the x86 primitives.
      #
      def _create_primitives

        #
        # Create the '_set_variable_location' primitive. The first param it the variable to place the current
        # blocks location value in.
        #
        _create_primitive( '_set_variable_location',
          ::Proc.new do | block, machine, variable |
            if( @safe_fpu_instructions.empty? )
              raise UnallowedPermutation
            end
            [
              "dw #{ "0x%04X" % [ @safe_fpu_instructions[ rand(@safe_fpu_instructions.length) ].unpack( 'v' ).first ] }",
              "mov #{machine.variable_value( 'temp' )}, esp",
              "fnstenv [ #{machine.variable_value( 'temp' )} - 12 ]",
              "pop #{machine.variable_value( variable )}"
            ]
          end,
          ::Proc.new do | block, machine, variable |
            if( @safe_fpu_instructions.empty? )
              raise UnallowedPermutation
            end
            [
              "dw #{ "0x%04X" % [ @safe_fpu_instructions[ rand(@safe_fpu_instructions.length) ].unpack( 'v' ).first ] }",
              "mov #{machine.variable_value( 'temp' )}, esp",
              "fnstenv [ #{machine.variable_value( 'temp' )} - 12 ]",
              "pop #{machine.variable_value( variable )}"
            ]
          end,
          ::Proc.new do | block, machine, variable |
            if( @safe_fpu_instructions.empty? )
              raise UnallowedPermutation
            end
            [
              "dw #{ "0x%04X" % [ @safe_fpu_instructions[ rand(@safe_fpu_instructions.length) ].unpack( 'v' ).first ] }",
              "push esp",
              "pop #{machine.variable_value( 'temp' )}",
              "fnstenv [ #{machine.variable_value( 'temp' )} - 12 ]",
              "pop #{machine.variable_value( variable )}"
            ]
          end,
          ::Proc.new do | block, machine, variable |
            if( @safe_fpu_instructions.empty? )
              raise UnallowedPermutation
            end
            [
              "dw #{ "0x%04X" % [ @safe_fpu_instructions[ rand(@safe_fpu_instructions.length) ].unpack( 'v' ).first ] }",
              "fnstenv [ esp - 12 ]",
              "pop #{machine.variable_value( variable )}"
            ]
          end,
          ::Proc.new do | block, machine, variable |
            [
              "call $+5",
              "pop #{machine.variable_value( variable )}",
              "push #{machine.block_offset( block ) + 5}",
              "pop #{machine.variable_value( 'temp' )}",
              "sub #{machine.variable_value( variable )}, #{machine.variable_value( 'temp' )}"
            ]
          end,
          ::Proc.new do | block, machine, variable |
            [
              "db 0xE8, 0xFF, 0xFF, 0xFF, 0xFF, 0xC0",
              "pop #{machine.variable_value( variable )}",
              "push #{machine.block_offset( block ) + 5}",
              "pop #{machine.variable_value( 'temp' )}",
              "sub #{machine.variable_value( variable )}, #{machine.variable_value( 'temp' )}"
            ]
          end
        )

        #
        # Create the 'loop' primitive. The first param it the counter variable which holds the number of
        # times to perform the loop. The second param it the destination block to loop to.
        #
        _create_primitive( 'loop',
          ::Proc.new do | block, machine, counter, destination |
            if( machine.variable_value( counter ) != Rex::Arch::X86::REG_NAMES32[ Rex::Arch::X86::ECX ] )
              # we raise and UndefinedPermutation exception to indicate that untill a valid register (ECX) is
              # chosen we simply can't render this. This lets the machine know we can still try to use this
              # permutation and at a later stage the requirements (counter==ecx) may be satisfied.
              raise UndefinedPermutation
            end
            offset = -( machine.block_offset( machine.block_next( block ) ) - machine.block_offset( destination ) )
            Rex::Arch::X86.loop( offset )
          end,
          ::Proc.new do | block, machine, counter, destination |
            offset = -( machine.block_offset( machine.block_next( block ) ) - machine.block_offset( destination ) )
            [
              "dec #{machine.variable_value( counter )}",
              "test #{machine.variable_value( counter )}, #{machine.variable_value( counter )}",
              # JNZ destination
              "db 0x0F, 0x85 dd #{ "0x%08X" % [ offset & 0xFFFFFFFF ] }"
            ]
          end
        )

        #
        # Create the 'xor' primitive. The first param it the variable to xor with the second param value which
        # can be either a variable, literal or block offset.
        #
        _create_primitive( 'xor',
          ::Proc.new do | block, machine, variable, value |
            [
              "xor #{machine.variable_value( variable )}, #{machine.resolve_value( value )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            # a ^ b == (a | b) & ~(a & b)
            [
              "mov #{machine.variable_value( 'temp' )}, #{machine.variable_value( variable )}",
              "or #{machine.variable_value( 'temp' )}, #{machine.resolve_value( value )}",
              "and #{machine.variable_value( variable )}, #{machine.resolve_value( value )}",
              "not #{machine.variable_value( variable )}",
              "and #{machine.variable_value( variable )}, #{machine.variable_value( 'temp' )}"
            ]
          end
        )

        #
        # Create the 'goto' primitive. The first param is a destination block to jump to.
        #
        _create_primitive( 'goto',
          ::Proc.new do | block, machine, destination |
            offset = -( machine.block_offset( machine.block_next( block ) ) - machine.block_offset( destination ) )
            if( ( offset > 0 and offset > 127 ) or ( offset < 0 and offset < -127 ) )
              raise UnallowedPermutation
            end
            [
              # short relative jump
              "db 0xEB db #{ "0x%02X" % [ offset & 0xFF ] }"
            ]
          end,
          ::Proc.new do | block, machine, destination |
            offset = -( machine.block_offset( machine.block_next( block ) ) - machine.block_offset( destination ) )
            [
              # near relative jump
              "db 0xE9 dd #{ "0x%08X" % [ offset & 0xFFFFFFFF ] }"
            ]
          end
        )

        #
        # Create the 'add' primitive. The first param it the variable which will be added to the second
        # param, which may either be a literal number value, a variables assigned register or a block
        # name, in which case the block offset will be used.
        #
        _create_primitive( 'add',
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            [
              "lea #{machine.variable_value( variable )}, [ #{machine.variable_value( variable )} + #{machine.resolve_value( value )} ]"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            [
              "push #{machine.resolve_value( value )}",
              "add #{machine.variable_value( variable )}, [esp]",
              "pop #{machine.variable_value( 'temp' )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            [
              "add #{machine.variable_value( variable )}, #{machine.resolve_value( value )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            [
              "sub #{machine.variable_value( variable )}, #{ "0x%08X" % [ ~(machine.resolve_value( value ) - 1) & 0xFFFFFFFF ] }"
            ]
          end
          # ::Proc.new do | block, machine, variable, value |
            # if( machine.variable_exist?( value ) )
            #	raise UnallowedPermutation
            # end
            # [
              # "push #{ "0x%08X" % [ ~(machine.resolve_value( value ) - 1) & 0xFFFFFFFF ] }",
              # "pop #{machine.variable_value( 'temp' )}",
              # "not #{machine.variable_value( 'temp' )}",
              # "add #{machine.variable_value( variable )}, #{machine.variable_value( 'temp' )}"
            # ]
          # end,
          # ::Proc.new do | block, machine, variable, value |
            # if( machine.variable_exist?( value ) )
            #	raise UnallowedPermutation
            # end
            # [
              # "xor #{machine.variable_value( 'temp' )}, #{machine.variable_value( 'temp' )}",
              # "mov #{machine.variable_value( 'temp', 16 )}, #{ "0x%04X" % [ ~(machine.resolve_value( value ) - 1) & 0xFFFF ] }",
              # "not #{machine.variable_value( 'temp', 16 )}",
              # "add #{machine.variable_value( variable )}, #{machine.variable_value( 'temp' )}"
            # ]
          # end,
        )

        #
        # Create the 'set' primitive. The first param it the variable which will be set. the second
        # param is the value to set the variable to (a variable, block or literal).
        #
        _create_primitive( 'set',
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            [
              "push #{ "0x%08X" % [ ~machine.resolve_value( value ) & 0xFFFFFFFF ] }",
              "pop #{machine.variable_value( variable )}",
              "not #{machine.variable_value( variable )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            if( machine.resolve_value( value, WORD ) > 0xFFFF )
              raise UndefinedPermutation
            end
            [
              "xor #{machine.variable_value( variable )}, #{machine.variable_value( variable )}",
              "mov #{machine.variable_value( variable, WORD )}, #{ "0x%04X" % [ ~machine.resolve_value( value, WORD ) & 0xFFFF ] }",
              "not #{machine.variable_value( variable, WORD )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            [
              "push #{machine.resolve_value( value )}",
              "pop #{machine.variable_value( variable )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            [
              "mov #{machine.variable_value( variable )}, #{machine.resolve_value( value )}"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            if( machine.resolve_value( value, WORD ) > 0xFFFF )
              raise UndefinedPermutation
            end
            [
              "xor #{machine.variable_value( variable )}, #{machine.variable_value( variable )}",
              "mov #{machine.variable_value( variable, WORD )}, #{ "0x%04X" % [ machine.resolve_value( value, WORD ) & 0xFFFF ] }"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            dword = machine.make_safe_dword( machine.resolve_value( value ) )
            [
              "mov #{machine.variable_value( variable )}, #{ "0x%08X" % [ dword ] }",
              "sub #{machine.variable_value( variable )}, #{ "0x%08X" % [ dword - machine.resolve_value( value ) ] }"
            ]
          end,
          ::Proc.new do | block, machine, variable, value |
            if( machine.variable_exist?( value ) )
              raise UnallowedPermutation
            end
            dword = machine.make_safe_dword( machine.resolve_value( value ) )
            [
              "mov #{machine.variable_value( variable )}, #{ "0x%08X" % [ dword - machine.resolve_value( value ) ] }",
              "add #{machine.variable_value( variable )}, #{ "0x%08X" % [ ~dword & 0xFFFFFFFF ] }",
              "not #{machine.variable_value( variable )}"
            ]
          end
        )

        #
        # Create the 'load' primitive. The first param it the variable which will be set. The second
        # param is the value (either a variable or literal) to load from. the third param is the size
        # of the load operation, either DWORD, WORD or BYTE.
        #
        _create_primitive( 'load',
          ::Proc.new do | block, machine, variable, value, size |
            result = nil
            if( size == Rex::Poly::Machine::DWORD )
              result = [ "mov #{machine.variable_value( variable )}, [#{machine.resolve_value( value )}]" ]
            elsif( size == Rex::Poly::Machine::WORD )
              result = [ "movzx #{machine.variable_value( variable )}, word [#{machine.resolve_value( value )}]" ]
            elsif( size == Rex::Poly::Machine::BYTE )
              result = [ "movzx #{machine.variable_value( variable )}, byte [#{machine.resolve_value( value )}]" ]
            else
              raise InvalidPermutation
            end
            result
          end,
          ::Proc.new do | block, machine, variable, value, size |
            result = nil
            if( size == Rex::Poly::Machine::DWORD )
              # we raise and UnallowedPermutation here as this permutation should only satisfy requests for
              # sizes of WORD or BYTE, any DWORD requests will be satisfied by the above permutation (otherwise
              # we would just be duplicating a 'mov dest, [src]' sequence which is the same as above.
              raise UnallowedPermutation
            elsif( size == Rex::Poly::Machine::WORD )
              result = [
                "mov #{machine.variable_value( variable )}, [#{machine.resolve_value( value )}]",
                "shl #{machine.variable_value( variable )}, 16",
                "shr #{machine.variable_value( variable )}, 16"
              ]
            elsif( size == Rex::Poly::Machine::BYTE )
              result = [
                "mov #{machine.variable_value( variable )}, [#{machine.resolve_value( value )}]",
                "shl #{machine.variable_value( variable )}, 24",
                "shr #{machine.variable_value( variable )}, 24"
              ]
            else
              raise InvalidPermutation
            end
            result
          end,
          ::Proc.new do | block, machine, variable, value, size |
            result = nil
            if( size == Rex::Poly::Machine::DWORD )
              result = [
                "push [#{machine.resolve_value( value )}]",
                "pop #{machine.variable_value( variable )}"
              ]
            elsif( size == Rex::Poly::Machine::WORD )
              result = [
                "push [#{machine.resolve_value( value )}]",
                "pop #{machine.variable_value( variable )}",
                "shl #{machine.variable_value( variable )}, 16",
                "shr #{machine.variable_value( variable )}, 16"
              ]
            elsif( size == Rex::Poly::Machine::BYTE )
              result = [
                "push [#{machine.resolve_value( value )}]",
                "pop #{machine.variable_value( variable )}",
                "shl #{machine.variable_value( variable )}, 24",
                "shr #{machine.variable_value( variable )}, 24"
              ]
            else
              raise InvalidPermutation
            end
            result
          end
        )

        #
        # Create the 'store' primitive.
        #
        _create_primitive( 'store',
          ::Proc.new do | block, machine, variable, value, size |
            result = nil
            if( size == Rex::Poly::Machine::DWORD )
              result = [ "mov [#{machine.variable_value( variable )}], #{machine.resolve_value( value )}" ]
            elsif( size == Rex::Poly::Machine::WORD )
              result = [ "mov word [#{machine.variable_value( variable )}], #{machine.resolve_value( value, WORD )}" ]
            elsif( size == Rex::Poly::Machine::BYTE )
              if( machine.resolve_value( value, BYTE ).nil? )
                # so long as we cant resolve the variable to an 8bit register value (AL,BL,CL,DL) we must raise
                # an UndefinedPermutation exception (this will happen when the variable has been assigned to ESI,
                # EDI, EBP or ESP which dont have a low byte representation)
                raise UndefinedPermutation
              end
              result = [ "mov byte [#{machine.variable_value( variable )}], #{machine.resolve_value( value, BYTE )}" ]
            else
              raise InvalidPermutation
            end
            result
          end,
          ::Proc.new do | block, machine, variable, value, size |
            result = nil
            if( size == Rex::Poly::Machine::DWORD )
              result = [
                "push #{machine.resolve_value( value )}",
                "pop [#{machine.variable_value( variable )}]"
              ]
            elsif( size == Rex::Poly::Machine::WORD )
              result = [
                "push #{machine.resolve_value( value, WORD )}",
                "pop word [#{machine.variable_value( variable )}]"
              ]
            else
              # we can never do this permutation for BYTE size (or any other size)
              raise UnallowedPermutation
            end
            result
          end
        )
      end

    end

  end

end