Improve payload size generation script

Author: adfoster-r7

lib/msf/core/payload.rb

@@ -159,15 +159,25 @@ def staged?
   # This method returns an optional cached size value
   #
   def self.cached_size
-    csize = (const_defined?('CachedSize')) ? const_get('CachedSize') : nil
+    csize = const_defined?('CachedSize') ? const_get('CachedSize') : nil
+    if ancestors.include?(Msf::Payload::Stager)
+      csize_overrides = const_defined?('CachedSizeOverrides') ? const_get('CachedSizeOverrides') : {}
+      csize = csize_overrides.fetch(self.refname, csize)
+    end
+
     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 = const_defined?('CachedSize') ? const_get('CachedSize') : nil
+    if ancestors.include?(Msf::Payload::Stager)
+      csize_overrides = const_defined?('CachedSizeOverrides') ? const_get('CachedSizeOverrides') : {}
+      csize = csize_overrides.fetch(self.refname, csize)
+    end
+
     csize == :dynamic
   end
 

lib/msf/core/payload/java.rb

@@ -2,6 +2,9 @@
 
 module Msf::Payload::Java
 
+  # Mark the payload as dynamic as the generated JAR/zip files can differ in size depending on the host machine's zlib version
+  ForceDynamicCachedSize = true
+
   #
   # Used by stages; all java stages need to define +stage_class_files+ as an
   # array of .class files located in data/java/

lib/msf/core/payload/java/reverse_http.rb

@@ -15,6 +15,9 @@ module Payload::Java::ReverseHttp
   include Msf::Payload::UUID::Options
   include Msf::Payload::Java::PayloadOptions
 
+  # Mark the payload as dynamic as random length URIs are generated, and zip compression with a single char change can lead to size changes even if the original payload is the same length
+  ForceDynamicCachedSize = true
+
   #
   # Register Java reverse_http specific options
   #

lib/msf/core/payload/python.rb

@@ -1,6 +1,8 @@
 # -*- coding: binary -*-
 
 module Msf::Payload::Python
+  # Mark the payload as dynamic, as the zlib compression with a single char change can lead to size changes even if the original payload is the same length
+  ForceDynamicCachedSize = true
 
   #
   # Encode the given python command in base64 and wrap it with a stub

lib/msf/core/payload/python/meterpreter_loader.rb

@@ -10,6 +10,8 @@ module Msf
 ###
 
 module Payload::Python::MeterpreterLoader
+  # Mark the payload as dynamic, as random uuid values lead to differing zlib compressed payloads
+  ForceDynamicCachedSize = true
 
   include Msf::Payload::Python
   include Msf::Payload::UUID::Options

lib/msf/core/payload/windows/powershell.rb

@@ -9,6 +9,8 @@ module Msf
 ###
 
 module Payload::Windows::Powershell
+  # Mark the payload as dynamic as powershell scripts are randomized
+  ForceDynamicCachedSize = true
 
   def initialize(info = {})
     ret = super(info)

lib/msf/util/payload_cached_size.rb

@@ -17,8 +17,10 @@ class PayloadCachedSize
   OPTS = {
     'Format'      => 'raw',
     'Options'     => {
+      'VERBOSE' => false,
       'CPORT' => 4444,
       'LPORT' => 4444,
+      'RPORT' => 4444,
       'CMD' => '/bin/sh',
       'URL' => 'http://a.com',
       'PATH' => '/',
@@ -46,88 +48,196 @@ class PayloadCachedSize
 
   OPTS_IPV4 = {
     'LHOST' => '255.255.255.255',
+    'RHOST' => '255.255.255.255',
     'KHOST' => '255.255.255.255',
     'AHOST' => '255.255.255.255'
   }.freeze
 
   OPTS_IPV6 = {
     'LHOST' => 'ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff',
+    'RHOST' => 'ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff',
     'KHOST' => 'ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff',
     'AHOST' => 'ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff'
   }.freeze
 
-  # Insert a new CachedSize value into the text of a payload module
+  # Inserts or updates the CachedSize constant in 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]
+  # @param cached_size [String, Integer] The new value for CachedSize, which should be either an integer or the string ":dynamic"
+  # @return [String] The updated source code with the new CachedSize value
   def self.update_cache_constant(data, cached_size)
     data.
       gsub(/^\s*CachedSize\s*=\s*(\d+|:dynamic).*/, '').
       gsub(/^(module MetasploitModule)\s*\n/) do |m|
-        "#{m.strip}\n\n  CachedSize = #{cached_size}\n\n"
+        "#{m.strip}\n  CachedSize = #{cached_size}\n\n"
       end
   end
 
-  # Insert a new CachedSize value into a payload module file
+  # Inserts or updates the CachedSizeOverrides constant in the text of a payload module,
+  # removing any previous CachedSizeStages, # Other stager sizes, or CachedSizeOverrides lines.
+  #
+  # @param data [String] The source code of a payload module
+  # @param stages_with_sizes [Array<{:stage => Msf::Payload::Stager, :size => Integer}>] Array of hashes with :stage (an Msf::Payload::Stager instance) and :size (Integer)
+  # @return [String] The updated source code with the new CachedSizeOverrides value
+  def self.update_stage_sizes_constant(data, stages_with_sizes)
+    sizes = stages_with_sizes.sort_by { |stage_with_size| stage_with_size[:stage].refname }.map do |stage_with_size|
+      [stage_with_size[:stage].refname, stage_with_size[:size]]
+    end
+      data_without_other_stages = data.gsub(/^\s*CachedSizeOverrides\s*=.*\n/, '')
+    return data_without_other_stages if sizes.empty?
+
+    data_without_other_stages.gsub(/^\s*(CachedSize\s*=\s*(\d+|:dynamic))\s*\n/) do |m|
+      "  #{m.strip}\n  CachedSizeOverrides = {#{sizes.map { |(k, v)| %Q{"#{k}" => #{v}}}.join(', ')}}\n\n"
+    end
+  end
+
+  # Insert or update the 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
+  # @param cached_size [String, Integer] The new value for cached_size, which
+  #   should be either an integer or the string ":dynamic"
   # @return [void]
   def self.update_cached_size(mod, cached_size)
     mod_data = ""
 
-    ::File.open(mod.file_path, 'rb') do |fd|
+    file_path = mod.file_path
+
+    ::File.open(file_path, 'rb') do |fd|
       mod_data = fd.read(fd.stat.size)
     end
 
-    ::File.open(mod.file_path, 'wb') do |fd|
+    ::File.open(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
+  # Insert or update the CachedSize value into a payload module file
   #
   # @param mod [Msf::Payload] The class of the payload module to update
+  # @param stages_with_sizes [Array<{:stage => Msf::Payload::Stager, :size => Integer}>] Array of hashes with :stage (an Msf::Payload::Stager instance) and :size (Integer)
   # @return [void]
-  def self.update_module_cached_size(mod)
-    update_cached_size(mod, compute_cached_size(mod))
+  def self.update_stager_cached_sizes(mod, stages_with_sizes)
+    mod_data = ""
+
+    file_path = mod.file_path
+
+    ::File.open(file_path, 'rb') do |fd|
+      mod_data = fd.read(fd.stat.size)
+    end
+
+    ::File.open(file_path, 'wb') do |fd|
+      fd.write update_stage_sizes_constant( mod_data, stages_with_sizes)
+    end
+  end
+
+  # Updates the payload module specified with the current CachedSize
+  #
+  # @param framework [Msf::Framework] The Metasploit framework instance used for payload generation
+  # @param mod [Msf::Payload] The class of the payload module to update
+  # @return [String, Integer] The updated CachedSize value
+  def self.update_module_cached_size(framework, mod)
+    cached_size = compute_cached_size(framework, mod)
+    update_cached_size(mod, cached_size)
+    cached_size
+  end
+
+  # Updates the stager payload module with the most frequent CachedSize value and sets CachedSizeOverrides for other stages.
+  #
+  # @param framework [Msf::Framework] The Metasploit framework instance used for payload generation
+  # @param stages [Array<Msf::Payload>] Array of stager modules to update
+  # @return [Integer, String] The new CachedSize value set for the stager
+  def self.update_stager_module_cached_size(framework, stages)
+    stages_with_sizes = stages.map do |stage|
+      { stage: stage, size: compute_cached_size(framework, stage) }
+    end
+    most_frequent_cached_size = stages_with_sizes.map { |stage_with_size| stage_with_size[:size] }
+                            .select { |size| size.is_a?(Numeric) }.tally.sort_by(&:last).to_h.keys.last
+
+    new_size = most_frequent_cached_size || stages_with_sizes.first[:size]
+    other_sizes = stages_with_sizes.select { |stage_with_size| stage_with_size[:size] != new_size }
+
+    update_cached_size(stages.first, new_size)
+    update_stager_cached_sizes(stages.first, other_sizes)
+
+    new_size
   end
 
   # Calculates the CachedSize value for a payload module
   #
   # @param mod [Msf::Payload] The class of the payload module to update
-  # @return [Integer]
-  def self.compute_cached_size(mod)
-    return ":dynamic" if is_dynamic?(mod)
+  # @return [Integer, String]
+  def self.compute_cached_size(framework, mod)
+    return ":dynamic" if is_dynamic?(framework, mod)
 
-    mod.generate_simple(module_options(mod)).size
+    mod.replicant.generate_simple(module_options(mod)).bytesize
   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 [Integer] The number of iterations to use to
   #   verify that the size is static.
-  # @return [Integer]
-  def self.is_dynamic?(mod, generation_count=5)
+  # @return [Boolean]
+  def self.is_dynamic?(framework, mod, generation_count=10)
+    return true if mod.class.const_defined?('ForceDynamicCachedSize') && mod.class::ForceDynamicCachedSize
     opts = module_options(mod)
-    [*(1..generation_count)].map do |x|
-      mod.generate_simple(opts).size
-    end.uniq.length != 1
+    last_bytesize = nil
+    generation_count.times do
+      # Ensure a new module instance is created for each attempt, as some options are randomized on load - such as tmp file path names etc
+      new_mod = framework.payloads.create(mod.refname)
+      bytesize = new_mod.generate_simple(opts).bytesize
+      last_bytesize ||= bytesize
+      if last_bytesize != bytesize
+        return true
+      end
+    end
+
+    false
   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? && is_dynamic?(mod)
+  def self.is_cached_size_accurate?(framework, mod)
+    return true if mod.dynamic_size? && is_dynamic?(framework, mod)
     return false if mod.cached_size.nil?
 
-    mod.cached_size == mod.generate_simple(module_options(mod)).size
+    mod.cached_size == mod.replicant.generate_simple(module_options(mod)).bytesize
+  end
+
+  # Checks for errors or inconsistencies in the CachedSize value for a payload module.
+  # Returns nil if the cache is correct, or a string describing the error if not.
+  #
+  # @param framework [Msf::Framework] The Metasploit framework instance used for payload generation
+  # @param mod [Msf::Payload] The payload module to check
+  # @return [String, nil] Error message if there is a problem, or nil if the cache is correct
+  def self.cache_size_errors_for(framework, mod)
+    is_payload_size_different_on_each_generation = is_dynamic?(framework,mod)
+    module_marked_as_dynamic = mod.dynamic_size?
+    payload_cached_static_size = mod.cached_size
+
+    # Validate dynamic scenario
+    return if is_payload_size_different_on_each_generation && module_marked_as_dynamic
+
+    if is_payload_size_different_on_each_generation && !module_marked_as_dynamic
+      return 'Module generated different sizes for each generation attempt. CacheSize must be set to :dynamic'
+    end
+
+    if payload_cached_static_size.nil?
+      return 'Module missing CachedSize and not marked as dynamic'
+    end
+
+    payload_size_after_one_generation = mod.replicant.generate_simple(module_options(mod)).bytesize
+
+    # Validate static scenario
+    return if payload_cached_static_size == payload_size_after_one_generation
+
+    if payload_cached_static_size != payload_size_after_one_generation
+      return "Module marked as having size #{payload_cached_static_size} but after one generation was #{payload_size_after_one_generation}"
+    end
+
+    raise "unhandled scenario"
   end
 
   # Get a set of sane default options for the module so it can generate a

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

@@ -4,7 +4,7 @@
 ##
 
 module MetasploitModule
-  CachedSize = 164
+  CachedSize = 168
 
   include Msf::Payload::Single
   include Msf::Payload::Bsd

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

@@ -4,7 +4,7 @@
 ##
 
 module MetasploitModule
-  CachedSize = 128
+  CachedSize = 132
 
   include Msf::Payload::Single
   include Msf::Payload::Bsd

modules/payloads/singles/cmd/mainframe/apf_privesc_jcl.rb

@@ -20,6 +20,7 @@
 
 module MetasploitModule
   CachedSize = 3156
+
   include Msf::Payload::Single
   include Msf::Payload::Mainframe