Move the cache update logic into a utility class

lib/msf/util/payload_cached_size.rb

@@ -0,0 +1,59 @@
+# -*- coding: binary -*-
+###
+#
+#
+###
+
+module Msf
+module Util
+
+#
+# The class provides helper methods for verifying and updating the embedded CachedSize
+# constant within payload modules.
+#
+
+class PayloadCachedSize
+
+  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
+
+  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
+
+  def self.update_module_cached_size(mod)
+    update_cached_size(mod, compute_cached_size(mod))
+  end
+
+  def self.compute_cached_size(mod)
+    return :dynamic if is_dynamic?(mod)
+    return mod.new.size
+  end
+
+  def self.is_dynamic?(mod,generation_count=5)
+    [*(1..generation_count)].map{|x| mod.new.size}.uniq.length != 1
+  end
+
+  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

tools/update_payload_cached_sizes.rb

@@ -1,10 +1,6 @@
 #!/usr/bin/env ruby
 #
-# $Id$
-#
-# This script lists each exploit module by its compatible payloads
-#
-# $Revision$
+# This script updates the CachedSize constants in payload modules
 #
 
 msfbase = __FILE__
@@ -20,54 +16,14 @@ $:.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, val)
-  data = ''
-  File.open(mod.file_path, 'rb'){|fd| data = fd.read(fd.stat.size)}
-  data = data.gsub(/^\s*CachedSize\s*=\s*(\d+|:dynamic).*/, '')
-  data = data.gsub(/^(module Metasploit\d+)\s*\n/) {|m| "#{m.strip}\n\n  CachedSize = #{val}\n\n" }
-  File.open(mod.file_path, 'wb'){|fd| fd.write(data) }
-end
+require 'msf/util/payload_cached_size'
 
 # Initialize the simplified framework instance.
-$framework = Msf::Simple::Framework.create('DisableDatabase' => true)
+framework = Msf::Simple::Framework.create('DisableDatabase' => true)
 
-$framework.payloads.each_module do |name, mod|
-  gsize = mod.new.size
-
-  if is_dynamic_size?(mod) && ! mod.dynamic_size?
-    print_status("#{mod.file_path} has a dynamic size, updating cache...")
-    update_cache_size(mod, ":dynamic")
-    next
-  end
-
-  next if mod.dynamic_size?
-
-  if mod.cached_size.nil?
-    print_status("#{mod.file_path} has size #{gsize}, updating cache...")
-    update_cache_size(mod, gsize)
-  else
-    next if gsize == mod.cached_size
-    print_error("#{mod.file_path} has cached size #{mod.cached_size} but generated #{gsize}, updating cache...")
-    update_cache_size(mod, gsize)
-    next
-  end
+framework.payloads.each_module do |name, mod|
+  next if Msf::Util::PayloadCachedSize.is_cached_size_accurate?(mod)
+  $stdout.puts "[*] Updating the CacheSize for #{mod.file_path}..."
+  Msf::Util::PayloadCachedSize.update_module_cached_size(mod)
 end