rex-encoder 0.1.0

data/lib/rex/poly/machine.rb

@@ -0,0 +1,13 @@
+# -*- coding: binary -*-
+
+module Rex
+
+  module Poly
+
+    require 'metasm'
+    require 'rex/poly/machine/machine'
+    require 'rex/poly/machine/x86'
+
+  end
+
+end

data/lib/rex/poly/machine/machine.rb

@@ -0,0 +1,830 @@
+# -*- coding: binary -*-
+
+module Rex
+
+  module Poly
+
+    #
+    # A machine capable of creating a small blob of code in a metamorphic kind of way.
+    # Note: this is designed to perform an exhaustive search for a solution and can be
+    # slow. If you need a speedier option, the origional Rex::Polly::Block stuff is a
+    # better choice.
+    #
+    class Machine
+
+      QWORD = 8
+      DWORD = 4
+      WORD  = 2
+      BYTE  = 1
+
+      #
+      # A Permutation!
+      #
+      class Permutation
+
+        attr_accessor :active, :offset
+
+        attr_reader :name, :primitive, :length, :args
+
+        #
+        # Create a new permutation object.
+        #
+        def initialize( name, primitive, machine, source, args=nil )
+          @name      = name
+          @primitive = primitive
+          @machine   = machine
+          @source    = source
+          @args      = args
+          @active    = false
+          @valid     = true
+          @length    = 0
+          @offset    = 0
+          @children  = ::Array.new
+        end
+
+        #
+        # Add in a child permutation to this one. Used to build the permutation tree.
+        #
+        def add_child( child )
+          @children << child
+        end
+
+        #
+        # Does this permutation have children?
+        #
+        def has_children?
+          not @children.empty?
+        end
+
+        #
+        # Remove any existing children. Called by the machines generate function
+        # to build a fresh tree in case generate was previously called.
+        #
+        def remove_children
+          @children.clear
+        end
+
+        #
+        # Actully render this permutation into a raw buffer.
+        #
+        def render
+          raw   = ''
+          # Zero the length as we will be rendering the raw buffer and the length may change.
+          @length = 0
+          # If this permutation source is a Primitive/Procedure we can call it, otherwise we have a string
+          if( @source.kind_of?( Primitive ) or @source.kind_of?( ::Proc ) )
+            if( @source.kind_of?( Primitive ) )
+              raw = @source.call( @name, @machine, *@args )
+            elsif( @source.kind_of?( ::Proc ) )
+              raw = @source.call
+            end
+            # If the primitive/procedure returned an array, it is an array of assembly strings which we can assemble.
+            if( raw.kind_of?( ::Array ) )
+              lines = raw
+              raw   = ''
+              # itterate over each line of assembly
+              lines.each do | asm |
+                # parse the asm and substitute in any offset values specified...
+                offsets = asm.scan( /:([\S]+)_offset/ )
+                offsets.each do | name, |
+                  asm = asm.gsub( ":#{name}_offset", @machine.block_offset( name ).to_s )
+                end
+                # and substitute in and register values for any variables specified...
+                regs = asm.scan( /:([\S]+)_reg([\d]+)/ )
+                regs.each do | name, size |
+                  asm = asm.gsub( ":#{name}_reg#{size}", @machine.variable_value( name, size.to_i ) )
+                end
+                # assemble it into a raw blob
+                blob = @machine.assemble( asm )
+                #if( not @machine.is_valid?( blob ) )
+                #	p "#{name}(#{primitive}):#{asm} is invalid"
+                #end
+                raw << blob
+              end
+            end
+          else
+            # the source must just be a static string
+            raw = @source
+          end
+          # Update the length to reflect the new raw buffer
+          @length = raw.to_s.length
+          # As the temp variable is only assigned for the duration of a single permutation we
+          # can now release it if it was used in this permutation.
+          @machine.release_temp_variable
+          return raw.to_s
+        end
+
+        #
+        # Test if this permutation raw buffer is valid in this machine (e.g. against the badchar list).
+        #
+        def is_valid?
+          result = false
+          if( @valid )
+            begin
+              result = @machine.is_valid?( self.render )
+            rescue UnallowedPermutation
+              # This permutation is unallowed and can never be rendered so just mark it as
+              # not valid to skip it during future attempts.
+              @valid = false
+            rescue UndefinedPermutation
+              # allow an undefined permutation to fail validation but keep it marked
+              # as valid as it may be defined and passed validation later.
+            ensure
+              # Should a temporary variable have been assigned we can release it here.
+              @machine.release_temp_variable
+            end
+          end
+          return result
+        end
+
+        #
+        # Try to find a solution within the solution space by performing a depth first search
+        # into the permutation tree and backtracking when needed.
+        #
+        def solve
+          # Check to see if this permutation can make part of a valid solution
+          if( self.is_valid? )
+            # record this permutation as part of the final solution (the current machines register state is also saved here)
+            @machine.solution_push( self )
+            # If we have no children we are at the end of the tree and have a potential full solution.
+            if( not self.has_children? )
+              # We have a solution but doing a final pass to update offsets may introduce bad chars
+              # so we test for this and keep searching if this isnt a real solution after all.
+              if( not @machine.solution_is_valid? )
+                # remove this permutation and keep searching
+                @machine.solution_pop
+                return false
+              end
+              # Return true to unwind the recursive call as we have got a final solution.
+              return true
+            end
+            # Itterate over the children of this permutation (the perutations of the proceeding block).
+            @children.each do | child |
+              # Traverse into this child to keep trying to generate a solution...
+              if( child.solve )
+                # Keep returning true to unwind as we are done.
+                return true
+              end
+            end
+            # If we get here this permutation, origionally thought to be good for a solution, is not after all,
+            # so remove it from the machines final solution, restoring the register state aswell.
+            @machine.solution_pop
+          end
+          # No children can be made form part of the solution, return failure for this path in the tree.
+          return false
+        end
+
+      end
+
+      #
+      # A symbolic permutation to mark locations like the begining and end of a group of blocks.
+      # Used to calculate usefull offsets.
+      #
+      class SymbolicPermutation < Permutation
+        def initialize( name, machine, initial_offset=0 )
+          super( name, '', machine, '' )
+          # fudge the initial symbolic offset with a default (it gets patched correctly later),
+          # helps with the end symbolic block to not be 0 (as its a forward reference it really
+          # slows things down if we leave it 0)
+          @offset = initial_offset
+          # A symbolic block is allways active!
+          @active = true
+        end
+
+        #
+        # We block all attempts to set the active state of this permutation so as
+        # it is always true. This lets us always address the offset.
+        #
+        def active=( value )
+        end
+      end
+
+      #
+      # A primitive is a machine defined permutation which accepts some arguments when it is called.
+      #
+      class Primitive
+
+        #
+        # Initialize this primitive with its target source procedure and the machine it belongs to.
+        #
+        def initialize( source )
+          @source = source
+        end
+
+        #
+        # Call the primitives source procedure, passing in the arguments.
+        #
+        def call( name, machine, *args )
+          return @source.call( name, machine, *args )
+        end
+
+      end
+
+      #
+      #
+      #
+      class Block
+
+        #attr_accessor :next, :previous
+        attr_reader :name
+
+        def initialize( name )
+          @name         = name
+          @next         = nil
+          @previous     = nil
+          @permutations = ::Array.new
+        end
+
+        def shuffle
+          @permutations = @permutations.shuffle
+        end
+
+        def solve
+          @permutations.first.solve
+        end
+
+        def << ( permutation )
+          @permutations << permutation
+        end
+
+        def each
+          @permutations.each do | permutation |
+            yield permutation
+          end
+        end
+
+      end
+
+      #
+      # A class to hold a solution for a Rex::Poly::Machine problem.
+      #
+      class Solution
+
+        attr_reader :offset
+
+        def initialize
+          @permutations = ::Array.new
+          @reg_state    = ::Array.new
+          @offset       = 0
+        end
+
+        #
+        # Reset this solution to an empty state.
+        #
+        def reset
+          @offset = 0
+          @permutations.each do | permutation |
+            permutation.active = false
+            permutation.offset = 0
+          end
+          @permutations.clear
+          @reg_state.clear
+        end
+
+        #
+        # Push a new permutation onto this solutions permutations list and save the associated register/variables state
+        #
+        def push( permutation, reg_available, reg_consumed, variables )
+          permutation.active = true
+          permutation.offset = @offset
+          @offset += permutation.length
+          @permutations.push( permutation )
+          @reg_state.push( [ [].concat(reg_available), [].concat(reg_consumed), {}.merge(variables) ] )
+        end
+
+        #
+        # Pop off the last permutaion and register/variables state from this solution.
+        #
+        def pop
+          reg_available, reg_consumed, variables = @reg_state.pop
+          permutation = @permutations.pop
+          permutation.active = false
+          permutation.offset = 0
+          @offset -= permutation.length
+          return permutation, reg_available, reg_consumed, variables
+        end
+
+        #
+        # Render the final buffer.
+        #
+        def buffer
+          previous_offset = nil
+          count           = 0
+          # perform an N-pass fixup for offsets...
+          while( true ) do
+            # If we cant get the offsets fixed within a fixed ammount of tries we return
+            # nil to indicate failure and keep searching for a solution that will work.
+            if( count > 64 )
+              return nil
+            end
+            # Reset the solution offset so as to update it for this pass
+            @offset = 0
+            # perform a single pass to ensure we are using the correct offset values
+            @permutations.each do | permutation |
+              permutation.offset = @offset
+              # Note: calling render() can throw both UndefinedPermutation and UnallowedPermutation exceptions,
+              # however as we assume we only ever return the buffer once a final solution has been generated
+              # we should never have either of those exceptions thrown.
+              permutation.render
+              @offset += permutation.length
+            end
+            # If we have generated two consecutive passes which are the same length we can stop fixing up the offsets.
+            if( not previous_offset.nil? and @offset == previous_offset )
+              break
+            end
+            count +=1
+            previous_offset = @offset
+          end
+          # now a final pass to render the solution into the raw buffer
+          raw = ''
+          @permutations.each do | permutation |
+            #$stderr.puts "#{permutation.name} - #{ "0x%08X (%d)" % [ permutation.offset, permutation.length] } "
+            raw << permutation.render
+          end
+          return raw
+        end
+
+      end
+
+      #
+      # Create a new machine instance.
+      #
+      def initialize( badchars, cpu )
+        @badchars      = badchars
+        @cpu           = cpu
+
+        @reg_available = ::Array.new
+        @reg_consumed  = ::Array.new
+        @variables     = ::Hash.new
+        @blocks        = ::Hash.new
+        @primitives    = ::Hash.new
+        @solution      = Solution.new
+
+        _create_primitives
+
+        @blocks['begin'] = Block.new( 'begin' )
+        @blocks['begin'] << SymbolicPermutation.new( 'begin', self )
+
+        _create_variable( 'temp' )
+      end
+
+      #
+      # Overloaded by a subclass to return the maximum native general register size supported.
+      #
+      def native_size
+        nil
+      end
+
+      #
+      # Use METASM to assemble a line of asm using this machines current cpu.
+      #
+      def assemble( asm )
+        return Metasm::Shellcode.assemble( @cpu, asm ).encode_string
+      end
+
+      #
+      # Check if a data blob is valid against the badchar list (or perform any other validation here)
+      #
+      def is_valid?( data )
+        if( data.nil? )
+          return false
+        end
+        return Rex::Text.badchar_index( data, @badchars ).nil?
+      end
+
+      #
+      # Generate a 64 bit number whoes bytes are valid in this machine.
+      #
+      def make_safe_qword( number=nil )
+        return _make_safe_number( QWORD, number ) & 0xFFFFFFFFFFFFFFFF
+      end
+
+      #
+      # Generate a 32 bit number whoes bytes are valid in this machine.
+      #
+      def make_safe_dword( number=nil )
+        return _make_safe_number( DWORD, number ) & 0xFFFFFFFF
+      end
+
+      #
+      # Generate a 16 bit number whoes bytes are valid in this machine.
+      #
+      def make_safe_word( number=nil )
+        return _make_safe_number( WORD, number ) & 0xFFFF
+      end
+
+      #
+      # Generate a 8 bit number whoes bytes are valid in this machine.
+      #
+      def make_safe_byte( number=nil )
+        return _make_safe_number( BYTE, number ) & 0xFF
+      end
+
+      #
+      # Create a variable by name which will be assigned a register during generation. We can
+      # optionally assign a static register value to a variable if needed.
+      #
+      def create_variable( name, reg=nil )
+        # Sanity check we aren't trying to create one of the reserved variables.
+        if( name == 'temp' )
+          raise RuntimeError, "Unable to create variable, '#{name}' is a reserved variable name."
+        end
+        return _create_variable( name, reg )
+      end
+
+      #
+      # If the temp variable was assigned we release it.
+      #
+      def release_temp_variable
+        if( @variables['temp'] )
+          regnum = @variables['temp']
+          # Sanity check the temp variable was actually assigned (it may not have been if the last permutation didnot use it)
+          if( regnum )
+            # place the assigned register back in the available list for consumption later.
+            @reg_available.push( @reg_consumed.delete( regnum ) )
+            # unasign the temp vars register
+            @variables['temp'] = nil
+            return true
+          end
+        end
+        return false
+      end
+
+      #
+      # Resolve a variable name into its currently assigned register value.
+      #
+      def variable_value( name, size=nil )
+        # Sanity check we this variable has been created
+        if( not @variables.has_key?( name ) )
+          raise RuntimeError, "Unknown register '#{name}'."
+        end
+        # Pull out its current register value if it has been assigned one
+        regnum = @variables[ name ]
+        if( not regnum )
+          regnum = @reg_available.pop
+          if( not regnum )
+            raise RuntimeError, "Unable to assign variable '#{name}' a register value, none available."
+          end
+          # and add it to the consumed list so we can track it later
+          @reg_consumed << regnum
+          # and now assign the variable the register
+          @variables[ name ] = regnum
+        end
+        # resolve the register number int a string representation (e.g. 0 in x86 is EAX if size is 32)
+        return _register_value( regnum, size )
+      end
+
+      #
+      # Check this solution is still currently valid (as offsets change it may not be).
+      #
+      def solution_is_valid?
+        return self.is_valid?( @solution.buffer )
+      end
+
+      #
+      # As the solution advances we save state for each permutation step in the solution. This lets
+      # use rewind at a later stage if the solving algorithm wishes to perform some backtracking.
+      #
+      def solution_push( permutation )
+        @solution.push( permutation, @reg_available, @reg_consumed, @variables  )
+      end
+
+      #
+      # Backtrack one step in the solution and restore the register/variable state.
+      #
+      def solution_pop
+        permutation, @reg_available, @reg_consumed, @variables = @solution.pop
+
+        @reg_available.push( @reg_available.shift )
+      end
+
+      #
+      # Create a block by name and add in its list of permutations.
+      #
+      # XXX: this doesnt support the fuzzy order of block dependencies ala the origional rex::poly
+      def create_block( name, *permutation_sources )
+        # Sanity check we aren't trying to create one of the reserved symbolic blocks.
+        if( name == 'begin' or name == 'end' )
+          raise RuntimeError, "Unable to add block, '#{name}' is a reserved block name."
+        end
+        # If this is the first time this block is being created, create the block object to hold the permutation list
+        if( not @blocks[name] )
+          @blocks[name] = Block.new( name )
+        end
+        # Now create a new permutation object for every one supplied.
+        permutation_sources.each do | source |
+          @blocks[name] << Permutation.new( name, '', self, source )
+        end
+        return name
+      end
+
+      #
+      # Create a block which is based on a primitive defined by this machine.
+      #
+      def create_block_primitive( block_name, primitive_name, *args )
+        # Santiy check this primitive is actually available and is not an internal primitive (begins with an _).
+        if( not @primitives[primitive_name] or primitive_name[0] == "_" )
+          raise RuntimeError, "Unable to add block, Primitive '#{primitive_name}' is not available."
+        end
+        # Sanity check we aren't trying to create one of the reserved symbolic blocks.
+        if( block_name == 'begin' or block_name == 'end' )
+          raise RuntimeError, "Unable to add block, '#{block_name}' is a reserved block name."
+        end
+        return _create_block_primitive( block_name, primitive_name, *args )
+      end
+
+      #
+      # Get the offset for a blocks active permutation. This is easy for backward references as
+      # they will already have been rendered and their sizes known. For forward references we
+      # can't know in advance but the correct value can be known later once the final solution is
+      # available and a final pass to generate the raw buffer is made.
+      #
+      def block_offset( name )
+        if( name == 'end' )
+          return @solution.offset
+        elsif( @blocks[name] )
+          @blocks[name].each do | permutation |
+            if( permutation.active )
+              return permutation.offset
+            end
+          end
+        end
+        # If we are forward referencing a block it will be at least the current solutions offset +1
+        return @solution.offset + 1
+      end
+
+      #
+      # Does a given block exist?
+      #
+      def block_exist?( name )
+        return @blocks.include?( name )
+      end
+
+      #
+      # Does a given block exist?
+      #
+      def variable_exist?( name )
+        return @variables.include?( name )
+      end
+
+      # XXX: ambiguity between variable names and block name may introduce confusion!!! make them be unique.
+
+      #
+      # Resolve a given value into either a number literal, a block offset or
+      # a variables assigned register.
+      #
+      def resolve_value( value, size=nil )
+        if( block_exist?( value ) )
+          return block_offset( value )
+        elsif( variable_exist?( value ) )
+          return variable_value( value, size )
+        end
+        return value.to_i
+      end
+
+      #
+      # Get the block previous to the target block.
+      #
+      def block_previous( target_block )
+        previous_block = nil
+        @blocks.each_key do | current_block |
+          if( current_block == target_block )
+            return previous_block
+          end
+          previous_block = current_block
+        end
+        return nil
+      end
+
+      #
+      # Get the block next to the target block.
+      #
+      def block_next( target_block )
+        @blocks.each_key do | current_block |
+          if( block_previous( current_block ) == target_block )
+            return current_block
+          end
+        end
+        return nil
+      end
+
+      #
+      # Try to generate a solution.
+      #
+      def generate
+
+        if( @blocks.has_key?( 'end' ) )
+          @blocks.delete( 'end' )
+        end
+
+        @blocks['end'] = Block.new( 'end' )
+        @blocks['end'] << SymbolicPermutation.new( 'end', self, 1 )
+
+        # Mix up the permutation orders for each block and create the tree structure.
+        previous = ::Array.new
+        @blocks.each_value do | block |
+          # Shuffle the order of the blocks permutations.
+          block.shuffle
+          # create the tree by adding the current blocks permutations as children of the previous block.
+          current = ::Array.new
+          block.each do | permutation |
+            permutation.remove_children
+            previous.each do | prev |
+              prev.add_child( permutation )
+            end
+            current << permutation
+          end
+          previous = current
+        end
+
+        # Shuffle the order of the available registers
+        @reg_available = @reg_available.shuffle
+
+        # We must try every permutation of the register orders, so if we fail to
+        # generate a solution we rotate the available registers to try again with
+        # a different order. This ensures we perform and exhaustive search.
+        0.upto( @reg_available.length - 1 ) do
+
+          @solution.reset
+
+          # Start from the root node in the solution space and generate a
+          # solution by traversing the solution space's tree structure.
+          if( @blocks['begin'].solve )
+            # Return the solutions buffer (perform a last pass to fixup all offsets)...
+            return @solution.buffer
+          end
+
+          @reg_available.push( @reg_available.shift )
+        end
+
+        # :(
+        nil
+      end
+
+      #
+      # An UndefinedPermutation exception is raised when a permutation can't render yet
+      # as the conditions required are not yet satisfied.
+      #
+      class UndefinedPermutation < RuntimeError
+        def initialize( msg=nil )
+          super
+        end
+      end
+
+      #
+      # An UnallowedPermutation exception is raised when a permutation can't ever render
+      # as the conditions supplied are impossible to satisfy.
+      #
+      class UnallowedPermutation < RuntimeError
+        def initialize( msg=nil )
+          super
+        end
+      end
+
+      #
+      # An InvalidPermutation exception is raised when a permutation receives a invalid
+      # argument and cannot continue to render. This is a fatal exception.
+      #
+      class InvalidPermutation < RuntimeError
+        def initialize( msg=nil )
+          super
+        end
+      end
+
+      protected
+
+      #
+      # Overloaded by a subclass to resolve a register number into a suitable register
+      # name for the target architecture. E.g on x64 the register number 0 with size 64
+      # would resolve to RCX. Size is nil by default to indicate we want the default
+      # machine size, e.g. 32bit DWORD on x86 or 64bit QWORD on x64.
+      #
+      def _register_value( regnum, size=nil )
+        nil
+      end
+
+      #
+      # Perform the actual variable creation.
+      #
+      def _create_variable( name, reg=nil )
+        regnum = nil
+        # Sanity check this variable has not already been created.
+        if( @variables[name] )
+          raise RuntimeError, "Variable '#{name}' is already created."
+        end
+        # If a fixed register is being assigned to this variable then resolve it
+        if( reg )
+          # Resolve the register name into a register number
+          @reg_available.each do | num |
+            if( _register_value( num ) == reg.downcase )
+              regnum = num
+              break
+            end
+          end
+          # If an invalid register name was given or the chosen register is not available we must fail.
+          if( not regnum )
+            raise RuntimeError, "Register '#{reg}' is unknown or unavailable."
+          end
+          # Sanity check another variable isnt assigned this register
+          if( @variables.has_value?( regnum ) )
+            raise RuntimeError, "Register number '#{regnum}' is already consumed by variable '#{@variables[name]}'."
+          end
+          # Finally we consume the register chosen so we dont select it again later.
+          @reg_consumed << @reg_available.delete( regnum )
+        end
+        # Create the variable and assign it a register number (or nil if not yet assigned)
+        @variables[name] = regnum
+        return name
+      end
+
+      #
+      # Create a block which is based on a primitive defined by this machine.
+      #
+      def _create_block_primitive( block_name, primitive_name, *args )
+        # If this is the first time this block is being created, create the array to hold the permutation list
+        if( not @blocks[block_name] )
+          @blocks[block_name] = Block.new( block_name )
+        end
+        # Now create a new permutation object for every one supplied.
+        @primitives[primitive_name].each do | source |
+          @blocks[block_name] << Permutation.new( block_name, primitive_name, self, source, args )
+        end
+        return block_name
+      end
+
+      #
+      # Overloaded by a subclass to create any primitives available in this machine.
+      #
+      def _create_primitives
+        nil
+      end
+
+      #
+      # Rex::Poly::Machine::Primitive
+      #
+      def _create_primitive( name, *permutations )
+        # If this is the first time this primitive is being created, create the array to hold the permutation list
+        if( not @primitives[name] )
+          @primitives[name] = ::Array.new
+        end
+        # Add in the permutation object (Rex::Poly::Machine::Primitive) for every one supplied.
+        permutations.each do | permutation |
+          @primitives[name] << Primitive.new( permutation )
+        end
+      end
+
+      #
+      # Helper function to generate a number whoes byte representation is valid in this
+      # machine (does not contain any badchars for example). Optionally we can supply a
+      # number and the resulting addition/subtraction of this number against the newly
+      # generated value is also tested for validity. This helps in the assembly primitives
+      # which can use these values.
+      #
+      def _make_safe_number( bytes, number=nil )
+        format = ''
+        if( bytes == BYTE )
+          format = 'C'
+        elsif( bytes == WORD )
+          format = 'v'
+        elsif( bytes == DWORD )
+          format = 'V'
+        elsif( bytes == QWORD )
+          format = 'Q'
+        else
+          raise RuntimeError, "Invalid size '#{bytes}' used in _make_safe_number."
+        end
+
+        goodchars = (0..255).to_a
+
+        @badchars.unpack( 'C*' ).each do | b |
+          goodchars.delete( b.chr )
+        end
+
+        while( true ) do
+          value = 0
+
+          0.upto( bytes-1 ) do | i |
+            value |= ( (goodchars[ rand(goodchars.length) ] << i*8) & (0xFF << i*8) )
+          end
+
+          if( not is_valid?( [ value ].pack(format) ) or not is_valid?( [ ~value ].pack(format) ) )
+            redo
+          end
+
+          if( not number.nil? )
+            if(	not is_valid?( [ value + number ].pack(format) ) or not is_valid?( [ value - number ].pack(format) ) )
+              redo
+            end
+          end
+
+          break
+        end
+
+        return value
+      end
+
+    end
+
+  end
+
+end