This implements payload size caching, speeding up framework loads

lib/msf/core/payload.rb

@@ -159,6 +159,21 @@ class Payload < Msf::Module
     (@staged or payload_type == Type::Stager or payload_type == Type::Stage)
   end
 
+
+  #
+  # This method returns an optional cached size value
+  #
+  def self.cached_size
+    (const_defined?('CachedSize')) ? const_get('CachedSize') : nil
+  end
+
+  #
+  # This method returns an optional cached size value
+  #
+  def cached_size
+      self.class.cached_size
+  end
+
   #
   # Returns the payload's size.  If the payload is staged, the size of the
   # first stage is returned.
@@ -500,6 +515,12 @@ class Payload < Msf::Module
   #
   attr_accessor :assoc_exploit
 
+  #
+  # The amount of space available to the payload, which may be nil,
+  # indicating that the smallest possible payload should be used.
+  #
+  attr_accessor :available_space
+
 protected
 
   #

lib/msf/core/payload_set.rb

@@ -154,8 +154,7 @@ class PayloadSet < ModuleSet
           'type'  => op[5]['type']})
         new_keys.push combined
 
-        # Cache the payload's size
-        sizes[combined] = p.new.size
+        sizes[combined] = p.cached_size || p.new.size
       }
     }
 
@@ -236,7 +235,7 @@ class PayloadSet < ModuleSet
       next if (handler and not p.handler_klass.ancestors.include?(handler))
 
       # Check to see if the session classes match.
-      next if (session and p.session and not p.session.ancestors.include?(session))
+      next if (session and not p.session.ancestors.include?(session))
 
       # Check for matching payload types
       next if (payload_type and p.payload_type != payload_type)

tools/update_payload_cached_sizes.rb

@@ -0,0 +1,70 @@
+#!/usr/bin/env ruby
+#
+# $Id$
+#
+# This script lists each exploit module by its compatible payloads
+#
+# $Revision$
+#
+
+msfbase = __FILE__
+while File.symlink?(msfbase)
+  msfbase = File.expand_path(File.readlink(msfbase), File.dirname(msfbase))
+end
+
+$:.unshift(File.expand_path(File.join(File.dirname(msfbase), '..', 'lib')))
+require 'msfenv'
+
+$:.unshift(ENV['MSF_LOCAL_LIB']) if ENV['MSF_LOCAL_LIB']
+
+require 'rex'
+require 'msf/ui'
+require 'msf/base'
+
+
+def print_status(msg)
+  print_line "[*] #{msg}"
+end
+
+def print_error(msg)
+  print_line "[-] #{msg}"
+end
+
+def print_line(msg)
+  $stderr.puts msg
+end
+
+def is_dynamic_size?(mod)
+  [*(1..5)].map{|x| mod.new.size}.uniq.length != 1
+end
+
+def update_cache_size(mod)
+  data = ''
+  File.open(mod.file_path, 'rb'){|fd| data = fd.read(fd.stat.size)}
+  data = data.gsub(/^\s*CachedSize\s*=\s*\d+.*/, '')
+  data = data.gsub(/^(module Metasploit\d+)/) {|m| "#{m}\n  CachedSize = #{mod.new.size}\n" }
+  File.open(mod.file_path, 'wb'){|fd| fd.write(data) }
+end
+
+# Initialize the simplified framework instance.
+$framework = Msf::Simple::Framework.create('DisableDatabase' => true)
+
+$framework.payloads.each_module do |name, mod|
+  gsize = mod.new.size
+
+  if is_dynamic_size?(mod)
+    print_status("#{mod.file_path} has a dynamic size, skipping...")
+    next
+  end
+
+  if mod.cached_size.nil?
+    print_status("#{mod.file_path} has size #{gsize}, updating cache...")
+    update_cache_size(mod)
+  else
+    next if gsize == mod.cached_size
+    print_error("#{mod.file_path} has cached size #{mod.cached_size} but generated #{gsize}")
+    update_cache_size(mod)
+    next
+  end
+end
+