Land rapid7#4894: hmoore-r7 cache payload sizes on start

Avoid the hit of regenerating all of the static-size payloads when
loading the framework. This will facilitate conversion of payloads to
use metasm later.

Author: Brent Cook

lib/msf/core/payload.rb

@@ -159,6 +159,36 @@ def staged?
     (@staged or payload_type == Type::Stager or payload_type == Type::Stage)
   end
 
+  #
+  # This method returns an optional cached size value
+  #
+  def self.cached_size
+    csize = (const_defined?('CachedSize')) ? const_get('CachedSize') : nil
+    csize == :dynamic ? nil : csize
+  end
+
+  #
+  # This method returns whether the payload generates variable-sized output
+  #
+  def self.dynamic_size?
+    csize = (const_defined?('CachedSize')) ? const_get('CachedSize') : nil
+    csize == :dynamic
+  end
+
+  #
+  # This method returns an optional cached size value
+  #
+  def cached_size
+      self.class.cached_size
+  end
+
+  #
+  # This method returns whether the payload generates variable-sized output
+  #
+  def dynamic_size?
+      self.class.dynamic_size?
+  end
+
   #
   # Returns the payload's size.  If the payload is staged, the size of the
   # first stage is returned.

lib/msf/core/payload_set.rb

@@ -155,7 +155,7 @@ def recalculate
         new_keys.push combined
 
         # Cache the payload's size
-        sizes[combined] = p.new.size
+        sizes[combined] = p.cached_size || p.new.size
       }
     }
 

lib/msf/util/payload_cached_size.rb

@@ -0,0 +1,89 @@
+# -*- coding: binary -*-
+###
+#
+#
+###
+
+module Msf
+module Util
+
+#
+# The class provides helper methods for verifying and updating the embedded CachedSize
+# constant within payload modules.
+#
+
+class PayloadCachedSize
+
+  # Insert a new CachedSize value into the text of a payload module
+  #
+  # @param data [String] The source code of a payload module
+  # @param cached_size [String] The new value for cached_size, which
+  #   which should be either numeric or the string :dynamic
+  # @return [String]
+  def self.update_cache_constant(data, cached_size)
+    data.
+      gsub(/^\s*CachedSize\s*=\s*(\d+|:dynamic).*/, '').
+      gsub(/^(module Metasploit\d+)\s*\n/) do |m|
+        "#{m.strip}\n\n  CachedSize = #{cached_size}\n\n"
+      end
+  end
+
+  # Insert a new CachedSize value into a payload module file
+  #
+  # @param mod [Msf::Payload] The class of the payload module to update
+  # @param cached_size [String] The new value for cached_size, which
+  #   which should be either numeric or the string :dynamic
+  # @return [void]
+  def self.update_cached_size(mod, cached_size)
+    mod_data = ""
+
+    ::File.open(mod.file_path, 'rb') do |fd|
+      mod_data = fd.read(fd.stat.size)
+    end
+
+    ::File.open(mod.file_path, 'wb') do |fd|
+      fd.write update_cache_constant(mod_data, cached_size)
+    end
+  end
+
+  # Updates the payload module specified with the current CachedSize
+  #
+  # @param mod [Msf::Payload] The class of the payload module to update
+  # @return [void]
+  def self.update_module_cached_size(mod)
+    update_cached_size(mod, compute_cached_size(mod))
+  end
+
+  # Calculates the CachedSize value for a payload module
+  #
+  # @param mod [Msf::Payload] The class of the payload module to update
+  # @return [Fixnum]
+  def self.compute_cached_size(mod)
+    return ":dynamic" if is_dynamic?(mod)
+    return mod.new.size
+  end
+
+  # Determines whether a payload generates a static sized output
+  #
+  # @param mod [Msf::Payload] The class of the payload module to update
+  # @param generation_count [Fixnum] The number of iterations to use to
+  #   verify that the size is static.
+  # @return [Fixnum]
+  def self.is_dynamic?(mod,generation_count=5)
+    [*(1..generation_count)].map{|x| mod.new.size}.uniq.length != 1
+  end
+
+  # Determines whether a payload's CachedSize is up to date
+  #
+  # @param mod [Msf::Payload] The class of the payload module to update
+  # @return [Boolean]
+  def self.is_cached_size_accurate?(mod)
+    return true if mod.dynamic_size?
+    return false if mod.cached_size.nil?
+    mod.cached_size == mod.new.size
+  end
+
+end
+
+end
+end

modules/payloads/singles/aix/ppc/shell_bind_tcp.rb

@@ -10,6 +10,8 @@
 
 module Metasploit3
 
+  CachedSize = 264
+
   include Msf::Payload::Single
   include Msf::Payload::Aix
   include Msf::Sessions::CommandShellOptions

modules/payloads/singles/aix/ppc/shell_find_port.rb

@@ -10,6 +10,8 @@
 
 module Metasploit3
 
+  CachedSize = 220
+
   include Msf::Payload::Single
   include Msf::Payload::Aix
   include Msf::Sessions::CommandShellOptions

modules/payloads/singles/aix/ppc/shell_interact.rb

@@ -10,6 +10,8 @@
 
 module Metasploit3
 
+  CachedSize = 56
+
   include Msf::Payload::Single
   include Msf::Payload::Aix
   include Msf::Sessions::CommandShellOptions

modules/payloads/singles/aix/ppc/shell_reverse_tcp.rb

@@ -10,6 +10,8 @@
 
 module Metasploit3
 
+  CachedSize = 204
+
   include Msf::Payload::Single
   include Msf::Payload::Aix
   include Msf::Sessions::CommandShellOptions

modules/payloads/singles/bsd/sparc/shell_bind_tcp.rb

@@ -10,6 +10,8 @@
 
 module Metasploit3
 
+  CachedSize = 164
+
   include Msf::Payload::Single
   include Msf::Payload::Bsd
   include Msf::Sessions::CommandShellOptions

modules/payloads/singles/bsd/sparc/shell_reverse_tcp.rb

@@ -10,6 +10,8 @@
 
 module Metasploit3
 
+  CachedSize = 128
+
   include Msf::Payload::Single
   include Msf::Payload::Bsd
   include Msf::Sessions::CommandShellOptions

modules/payloads/singles/bsd/x86/exec.rb

@@ -17,6 +17,8 @@
 ###
 module Metasploit3
 
+  CachedSize = 107
+
   include Msf::Payload::Single
   include Msf::Payload::Bsd