Add profiling tools

Gemfile

@@ -20,6 +20,10 @@ group :development do
   gem 'pry-byebug'
   # module documentation
   gem 'octokit'
+  # memory profiling
+  gem 'memory_profiler'
+  # cpu profiling
+  gem 'ruby-prof'
   # Metasploit::Aggregator external session proxy
   # disabled during 2.5 transition until aggregator is available
   #gem 'metasploit-aggregator'

Gemfile.lock

@@ -194,6 +194,7 @@ GEM
     loofah (2.4.0)
       crass (~> 1.0.2)
       nokogiri (>= 1.5.9)
+    memory_profiler (0.9.14)
     metasm (1.0.4)
     metasploit-concern (2.0.5)
       activemodel (~> 4.2.6)
@@ -373,6 +374,7 @@ GEM
       ruby-progressbar (~> 1.7)
       unicode-display_width (>= 1.4.0, < 1.7)
     ruby-macho (2.2.0)
+    ruby-prof (1.3.0)
     ruby-progressbar (1.10.1)
     ruby-rc4 (0.1.5)
     ruby_smb (1.1.0)
@@ -427,6 +429,7 @@ PLATFORMS
 DEPENDENCIES
   factory_bot_rails
   fivemat
+  memory_profiler
   metasploit-framework!
   octokit
   pry-byebug
@@ -435,6 +438,7 @@ DEPENDENCIES
   rspec-rails
   rspec-rerun
   rubocop
+  ruby-prof
   simplecov (= 0.18.2)
   sqlite3 (~> 1.3.0)
   swagger-blocks

lib/metasploit/framework/command/console.rb

@@ -39,12 +39,56 @@ class Metasploit::Framework::Command::Console < Metasploit::Framework::Command::
     end
   end
 
+  def start_profiler
+    root_path = Pathname.new(Msf::Config::install_root)
+    tmp_directory = root_path.join('tmp')
+    timestamp = Time.now.strftime('%Y%m%d%H%M%S')
+
+    if ENV['METASPLOIT_CPU_PROFILE']
+      require 'ruby-prof'
+
+      dump_path = tmp_directory.join("#{timestamp}-cpu")
+      ::FileUtils.mkdir_p(dump_path)
+
+      RubyProf.start
+
+      at_exit do
+        puts "Generating CPU dump #{dump_path}"
+        result = RubyProf.stop
+        printer = RubyProf::MultiPrinter.new(result, [:flat, :graph_html, :tree, :stack])
+        printer.print(path: dump_path)
+
+        Rex::Compat.open_file(dump_path)
+      end
+    end
+
+    if ENV['METASPLOIT_MEMORY_PROFILE']
+      require 'memory_profiler'
+
+      ::FileUtils.mkdir_p(tmp_directory)
+      report_name = "#{timestamp}-memory.profile"
+      report_path = tmp_directory.join(report_name)
+
+      MemoryProfiler.start
+
+      at_exit do
+        puts "Generating memory report #{dump_path}"
+        report = MemoryProfiler.stop
+        report.pretty_print(to_file: report_path)
+
+        puts "Memory report saved to #{report_path}"
+        Rex::Compat.open_file(report_path)
+      end
+    end
+  end
+
   def start
     case parsed_options.options.subcommand
     when :version
       $stderr.puts "Framework Version: #{Metasploit::Framework::VERSION}"
     else
       spinner unless parsed_options.options.console.quiet
+      start_profiler
       driver.run
     end
   end

msfconsole

@@ -5,44 +5,10 @@
 # framework.
 #
 
-#
-# Standard Library
-#
-
 require 'pathname'
 
-if ENV['METASPLOIT_FRAMEWORK_PROFILE'] == 'true'
-  gem 'perftools.rb'
-  require 'perftools'
-
-  formatted_time = Time.now.strftime('%Y%m%d%H%M%S')
-  root = Pathname.new(__FILE__).parent
-  profile_pathname = root.join('tmp', 'profiles', 'msfconsole', formatted_time)
-
-  profile_pathname.parent.mkpath
-  PerfTools::CpuProfiler.start(profile_pathname.to_path)
-
-  at_exit {
-    PerfTools::CpuProfiler.stop
-
-    puts "Generating pdf"
-
-    pdf_path = "#{profile_pathname}.pdf"
-
-    if Bundler.clean_system("pprof.rb --pdf #{profile_pathname} > #{pdf_path}")
-      puts "PDF saved to #{pdf_path}"
-
-      Rex::Compat.open_file(pdf_path)
-    end
-  }
-end
-
-#
-# Project
-#
-
-# @see https://github.com/rails/rails/blob/v3.2.17/railties/lib/rails/generators/rails/app/templates/script/rails#L3-L5
 begin
+  # @see https://github.com/rails/rails/blob/v3.2.17/railties/lib/rails/generators/rails/app/templates/script/rails#L3-L5
   require Pathname.new(__FILE__).realpath.expand_path.parent.join('config', 'boot')
   require 'metasploit/framework/command/console'
   require 'msf/core/payload_generator'