Stage encoding is now SaveRegister aware

lib/msf/core/encoded_payload.rb

@@ -158,6 +158,18 @@ class EncodedPayload
           next
         end
 
+        # If the caller explictly requires register preservation, make sure
+        # that the module in question can handle it. This is mostly used by
+        # the stage encoder path.
+        if (reqs['ForceSaveRegisters'] and
+            reqs['EncoderOptions'] and
+            (reqs['EncoderOptions']['SaveRegisters'].to_s.length > 0) and
+            (! self.encoder.preserves_registers?))
+          wlog("#{pinst.refname}: Encoder #{encoder.refname} does not preserve registers and the caller needs #{reqs['EncoderOptions']['SaveRegisters']} preserved.",
+            'core', LEV_1)
+          next
+        end
+
         # Import the datastore from payload (and likely exploit by proxy)
         self.encoder.share_datastore(pinst.datastore)
 

lib/msf/core/encoder.rb

@@ -397,6 +397,27 @@ class Encoder < Module
     buf
   end
 
+  #
+  # Determines whether the encoder can preserve registers at all
+  #
+  def preserves_registers?
+    false
+  end
+
+  #
+  # A list of registers always modified by the encoder
+  #
+  def modified_registers
+    []
+  end
+
+  #
+  # Determines whether the encoder can preserve the stack frame
+  #
+  def preserves_stack?
+    false
+  end
+
 protected
 
   #

lib/msf/core/payload/stager.rb

@@ -16,6 +16,7 @@ module Msf::Payload::Stager
       [
         Msf::OptBool.new("EnableStageEncoding", [ false, "Encode the second stage payload", false ]),
         Msf::OptString.new("StageEncoder", [ false, "Encoder to use if EnableStageEncoding is set", nil ]),
+        Msf::OptString.new("StageEncoderSaveRegisters", [ false, "Additional registers to preserve in the staged payload if EnableStageEncoding is set", "" ])
       ], Msf::Payload::Stager)
 
   end
@@ -196,6 +197,16 @@ module Msf::Payload::Stager
     false
   end
 
+
+  #
+  # Takes an educated guess at the list of registers an encoded stage
+  # would need to preserve based on the Convention
+  def encode_stage_preserved_registers
+    module_info['Convention'].to_s.scan(/\bsock([a-z]{3,}+)\b/).
+      map {|reg| reg.first }.
+      join(" ")
+  end
+
   # Encodes the stage prior to transmission
   # @return [String] Encoded version of +stg+
   def encode_stage(stg)
@@ -207,14 +218,21 @@ module Msf::Payload::Stager
       stage_enc_mod = datastore["StageEncoder"]
     end
 
+    # Allow the user to specify additional registers to preserve
+    saved_registers = (
+      datastore['StageEncoderSaveRegisters'].to_s + " "
+      encode_stage_preserved_registers
+    ).strip
+
     # Generate an encoded version of the stage.  We tell the encoding system
-    # to save edi to ensure that it does not get clobbered.
+    # to save certain registers to ensure that it does not get clobbered.
     encp = Msf::EncodedPayload.create(
       self,
-      'Raw'           => stg,
-      'Encoder'       => stage_enc_mod,
-      'SaveRegisters' => ['edi'],
-      'ForceEncode'   => true)
+      'Raw'                => stg,
+      'Encoder'            => stage_enc_mod,
+      'EncoderOptions'     => { 'SaveRegisters' => saved_registers },
+      'ForceSaveRegisters' => true,
+      'ForceEncode'        => true)
     print_status("Encoded stage with #{encp.encoder.refname}")
 
     # If the encoding succeeded, use the encoded buffer.  Otherwise, fall

modules/encoders/x86/shikata_ga_nai.rb

@@ -37,9 +37,16 @@ class Metasploit3 < Msf::Encoder::XorAdditiveFeedback
   # Generates the shikata decoder stub.
   #
   def decoder_stub(state)
+
     # If the decoder stub has not already been generated for this state, do
     # it now.  The decoder stub method may be called more than once.
     if (state.decoder_stub == nil)
+
+      # Sanity check that saved_registers doesn't overlap with modified_registers
+      if (modified_registers & saved_registers).length > 0
+        raise BadGenerateError
+      end
+
       # Shikata will only cut off the last 1-4 bytes of it's own end
       # depending on the alignment of the original buffer
       cutoff = 4 - (state.buf.length & 3)
@@ -61,6 +68,33 @@ class Metasploit3 < Msf::Encoder::XorAdditiveFeedback
     state.decoder_stub
   end
 
+  # Indicate that this module can preserve some registers
+  def preserves_registers?
+    true
+  end
+
+  # A list of registers always touched by this encoder
+  def modified_registers
+    # ESP is assumed and is handled through preserves_stack?
+    [
+      # The counter register is hardcoded
+      Rex::Arch::X86::ECX,
+      # These are modified by div and mul operations
+      Rex::Arch::X86::EAX, Rex::Arch::X86::EDX
+    ]
+  end
+
+  # Always preserve these registers in our block generation
+  def preserved_registers
+    ([
+      # Never modify our stack pointer
+      Rex::Arch::X86::ESP,
+      # Never modify our counter register
+      Rex::Arch::X86::ECX
+      # Never modify user specified registers
+    ] + saved_registers).uniq
+  end
+
 protected
 
   #
@@ -251,15 +285,31 @@ protected
     loop_inst.depends_on(loop_block)
 
     begin
-      # Generate a permutation saving the ECX and ESP registers
-      loop_inst.generate([
-        Rex::Arch::X86::ESP,
-        Rex::Arch::X86::ECX ], nil, state.badchars)
+      # Generate a permutation saving the ECX, ESP, and user defined registers
+      loop_inst.generate(preserved_registers, nil, state.badchars)
     rescue RuntimeError => e
       raise EncodingError
     end
   end
 
+  # Exploit or user-supplied list of registes to preserve
+  def saved_registers
+    saved = []
+
+    # Process any specified registers
+    datastore['SaveRegisters'].to_s.
+      strip.split(/[,\s]/).
+      map    {|reg| reg.to_s.strip.upcase }.
+      select {|reg| reg.length > 0        }.
+      uniq.each do |reg|
+        if Rex::Arch::X86.const_defined?(reg.intern)
+          saved << Rex::Arch::X86.const_get(reg.intern)
+        end
+    end
+
+    saved
+  end
+
   def sub_immediate(regnum, imm)
     return "" if imm.nil? or imm == 0
     if imm > 255 or imm < -255