CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
# frozen_string_literal: true
2

3
require 'pathname'
4
require 'tmpdir'
5

6
module Metasploit
7
  module Framework
8
    module Profiler
9
      class << self
10
        def start
11
          return unless record_global_cpu? || record_global_memory?
12
          raise 'Cannot profile memory and cpu at the same time' if record_global_cpu? && record_global_memory?
13

14
          if record_global_cpu?
15
            require 'ruby-prof'
16

17
            results_path = tmp_cpu_results_path
18
            profile = RubyProf::Profile.new
19
            profile.start
20

21
            at_exit do
22
              result = profile.stop
23
              save_cpu_result(result, path: results_path)
24
            end
25
          end
26

27
          if record_global_memory?
28
            require 'memory_profiler'
29

30
            results_path = tmp_memory_results_path
31
            profile = MemoryProfiler
32
            profile.start
33

34
            at_exit do
35
              puts "Generating memory dump #{results_path}"
36
              result = profile.stop
37
              save_memory_result(result, path: results_path)
38
            end
39
          end
40
        end
41

42
        def record_cpu
43
          require 'ruby-prof'
44

45
          results_path = tmp_cpu_results_path
46
          profile = RubyProf::Profile.new
47
          profile.start
48

49
          yield
50

51
          result = profile.stop
52
          save_cpu_result(result, path: results_path)
53
        end
54

55
        def record_memory
56
          raise 'Cannot mix global memory recording and localised memory recording' if record_global_memory?
57

58
          require 'memory_profiler'
59

60
          results_path = tmp_memory_results_path
61
          profile = MemoryProfiler
62
          profile.start
63

64
          yield
65

66
          result = profile.stop
67
          save_memory_result(result, path: results_path)
68
        end
69

70
        private
71

72
        def record_global_cpu?
73
          ENV['METASPLOIT_CPU_PROFILE']
74
        end
75

76
        def record_global_memory?
77
          ENV['METASPLOIT_MEMORY_PROFILE']
78
        end
79

80
        def tmp_path_for(name:)
81
          tmp_directory = Dir.mktmpdir("msf-profile-#{Time.now.strftime('%Y%m%d%H%M%S')}")
82
          Pathname.new(tmp_directory).join(name)
83
        end
84

85
        def tmp_cpu_results_path
86
          path = tmp_path_for(name: 'cpu')
87
          ::FileUtils.mkdir_p(path)
88
          path
89
        end
90

91
        def tmp_memory_results_path
92
          tmp_path_for(name: 'memory')
93
        end
94

95
        def save_cpu_result(result, path:)
96
          require 'rex/compat'
97

98
          puts "Generating CPU dump #{path}"
99

100
          printer = RubyProf::MultiPrinter.new(result, %i[flat graph_html tree stack])
101
          printer.print(path: path)
102

103
          Rex::Compat.open_file(path)
104
        end
105

106
        def save_memory_result(result, path:)
107
          require 'rex/compat'
108

109
          result.pretty_print(to_file: path)
110
          Rex::Compat.open_file(path)
111
        end
112
      end
113
    end
114
  end
115
end
116

117