diff --git a/lib/msf/util/payload_cached_size.rb b/lib/msf/util/payload_cached_size.rb new file mode 100644 index 0000000000..3479b0e883 --- /dev/null +++ b/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 diff --git a/tools/update_payload_cached_sizes.rb b/tools/update_payload_cached_sizes.rb index c398dd9c6a..950d519669 100755 --- a/tools/update_payload_cached_sizes.rb +++ b/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