1
#!/usr/bin/env ruby
2

3
##
4
# This module requires Metasploit: https://metasploit.com/download
5
# Current source: https://github.com/rapid7/metasploit-framework
6
##
7

8
#
9
# This script updates the CachedSize constants in payload modules
10
#
11

12
msfbase = __FILE__
13
msfbase = File.expand_path(File.readlink(msfbase), File.dirname(msfbase)) while File.symlink?(msfbase)
14

15
$LOAD_PATH.unshift(File.expand_path(File.join(File.dirname(msfbase), '..', '..', 'lib')))
16
require 'msfenv'
17

18
$LOAD_PATH.unshift(ENV['MSF_LOCAL_LIB']) if ENV['MSF_LOCAL_LIB']
19

20
gem 'rex-text'
21
require 'rex'
22

23
class StatusReporter
24
  CLEAR_LINE = "\r\e[2K\r".freeze
25
  private_constant :CLEAR_LINE
26

27
  def print_progress(s)
28
    return if s.nil?
29

30
    @last_progress = s
31
    print CLEAR_LINE
32
    print info(s)
33
    $stdout.flush
34
  end
35

36
  def print_info(s)
37
    print CLEAR_LINE
38
    puts info(s)
39
    print_progress(@last_progress)
40
  end
41

42
  def print_error(s)
43
    $stderr.print CLEAR_LINE
44
    $stderr.puts error(s)
45
    print_progress(@last_progress)
46
  end
47

48
  def finish
49
    print CLEAR_LINE
50
    puts
51
  end
52

53
  private
54

55
  def info(s)
56
    "\e[1;36m[*]\e[0m #{s}"
57
  end
58

59
  def error(s)
60
    "\e[1;33m[!]\e[0m #{s}"
61
  end
62
end
63

64
# Initialize the simplified framework instance.
65
framework = Msf::Simple::Framework.create('DisableDatabase' => true)
66
exceptions = []
67
reporter = StatusReporter.new
68
current_payload = 0
69

70
# Currently the cached size is stored on stagers, but multiple stages can be associated with one stager
71
# Maps the stager to the available stages
72
stagers_to_stages = Hash.new { |hash, key| hash[key] = [] }
73

74
modules = []
75
framework.payloads.each_module do |name, mod|
76
  modules << [name, mod]
77
end
78

79
total_payloads = modules.length
80
modules.each do |name, mod|
81
  next if name =~ /generic/
82

83
  current_payload += 1
84
  reporter.print_progress "Updating single (#{current_payload}/#{total_payloads}) #{name}..."
85
  mod_inst = framework.payloads.create(name)
86

87
  next if mod_inst.is_a?(Msf::Payload::Adapter)
88

89
  mod_dependencies = mod_inst.dependencies
90
  missing_dependencies = mod_dependencies.reject(&:available?)
91
  if missing_dependencies.any?
92
    reporter.print_error "Cannot update payload size for #{name} - missing dependencies: #{missing_dependencies.join(',')}"
93
    next
94
  end
95

96
  if mod_inst.is_a?(Msf::Payload::Stager)
97
    stagers_to_stages[mod_inst.file_path] << mod_inst
98
    next
99
  end
100

101
  current_size = mod.dynamic_size? ? ':dynamic' : mod.cached_size
102
  new_size = Msf::Util::PayloadCachedSize.update_module_cached_size(framework, mod_inst)
103
  if current_size != new_size
104
    reporter.print_info "Single Updated: #{name} CacheSize on disk at #{mod.file_path} from #{current_size} to #{new_size}..."
105
  end
106
rescue StandardError => e
107
  reporter.print_error "Caught Error while updating #{name}:\n#{e}\n#{e.backtrace.map { |line| "\t#{line}" }.join("\n")}"
108
  exceptions << [ e, name ]
109
end
110

111
# Update the metadata on the stager module associated with stages
112
stager_count = 0
113
stagers_to_stages.each_value do |stages|
114
  stager_count += 1
115
  mod = stages.first
116
  reporter.print_progress("Updating stager (#{stager_count}/#{stagers_to_stages.length}) #{stages.first.refname}...")
117
  if mod.dynamic_size?
118
    current_size = ':dynamic'
119
  else
120
    current_size = mod.class.const_defined?('CachedSize') ? mod.class.const_get('CachedSize') : nil
121
  end
122
  new_size = Msf::Util::PayloadCachedSize.update_stager_module_cached_size(framework, stages)
123
  if current_size != new_size
124
    reporter.print_info "Stager Updated: #{mod.refname} CacheSize on disk at #{mod.file_path} from #{current_size} to #{new_size}..."
125
  end
126
rescue StandardError => e
127
  reporter.print_error "Caught Error while updating #{mod.refname}:\n#{e}\n#{e.backtrace.map { |line| "\t#{line}" }.join("\n")}"
128
  exceptions << [ e, name ]
129
end
130

131
reporter.finish
132

133
exit(1) unless exceptions.empty?
134

135