Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
require 'msfdb_helpers/db_interface'
2

3
module MsfdbHelpers
4
  class PgCtl < DbInterface
5

6
    def initialize(db_path:, options:, localconf:, db_conf:)
7
      @db = db_path
8
      @options = options
9
      @localconf = localconf
10
      @db_conf = db_conf
11
      @socket_directory = db_path
12
      super(options)
13
    end
14

15
    def init(msf_pass, msftest_pass)
16
      puts "Creating database at #{@db}"
17
      Dir.mkdir(@db)
18
      run_cmd("initdb --auth-host=trust --auth-local=trust -E UTF8 #{@db.shellescape}")
19

20
      File.open("#{@db}/postgresql.conf", 'a') do |f|
21
        f.puts "port = #{@options[:db_port]}"
22
      end
23

24
      # Try creating a test file at {Dir.tmpdir},
25
      # Else fallback to creation at @{db}
26
      # Else fail with error.
27
      if test_executable_file("#{Dir.tmpdir}")
28
        @socket_directory = Dir.tmpdir
29
      elsif test_executable_file("#{@db}")
30
        @socket_directory = @db
31
      else
32
        print_error("Attempt to create DB socket file at Temporary Directory and `~/.msf4/db` failed. Possibly because they are mounted with NOEXEC flags. Database initialization failed.")
33
      end
34

35
      start
36

37
      create_db_users(msf_pass, msftest_pass)
38

39
      write_db_client_auth_config
40
      restart
41
    end
42

43
    # Creates and attempts to execute a testfile in the specified directory,
44
    # to determine if it is mounted with NOEXEC flags.
45
    def test_executable_file(path)
46
      begin
47
        file_name = File.join(path, 'msfdb_testfile')
48
        File.open(file_name, 'w') do |f|
49
          f.puts "#!/bin/bash\necho exec"
50
        end
51
        File.chmod(0744, file_name)
52

53
        if run_cmd(file_name)
54
          File.open("#{@db}/postgresql.conf", 'a') do |f|
55
            f.puts "unix_socket_directories = \'#{path}\'"
56
          end
57
          puts "Creating db socket file at #{path}"
58
        end
59
        return true
60

61
      rescue => e
62
        return false
63

64
      ensure
65
        begin
66
          File.delete(file_name)
67
        rescue
68
          print_error("Unable to delete test file #{file_name}")
69
        end
70
      end
71

72
    end
73

74
    def delete
75
      if exists?
76
        stop
77

78
        if @options[:delete_existing_data]
79
          puts "Deleting all data at #{@db}"
80
          FileUtils.rm_rf(@db)
81
        end
82

83
        if @options[:delete_existing_data]
84
          FileUtils.rm_r(@db_conf, force: true)
85
        end
86
      else
87
        puts "No data at #{@db}, doing nothing"
88
      end
89
    end
90

91
    def start
92
      if status == DatabaseStatus::RUNNING
93
        puts "Database already started at #{@db}"
94
        return true
95
      end
96

97
      print "Starting database at #{@db}..."
98
      pg_ctl_spawn_cmd = "pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} -l #{@db.shellescape}/log start &"
99
      puts "spawn_cmd: #{pg_ctl_spawn_cmd}" if @options[:debug]
100
      pg_ctl_pid = Process.spawn(pg_ctl_spawn_cmd)
101
      Process.detach(pg_ctl_pid)
102
      is_database_running = retry_until_truthy(timeout: 60) do
103
        status == DatabaseStatus::RUNNING
104
      end
105

106
      if is_database_running
107
        puts 'success'.green.bold.to_s
108
        true
109
      else
110
        begin
111
          Process.kill(:KILL, pg_ctl_pid)
112
        rescue => e
113
          puts "Failed to kill pg_ctl_pid=#{pg_ctl_pid} - #{e.class} #{e.message}" if @options[:debug]
114
        end
115
        puts 'failed'.red.bold.to_s
116
        false
117
      end
118
    end
119

120
    def stop
121
      if status == DatabaseStatus::RUNNING
122
        puts "Stopping database at #{@db}"
123
        run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} stop")
124
      else
125
        puts "Database is no longer running at #{@db}"
126
      end
127
    end
128

129
    def restart
130
      stop
131
      start
132
    end
133

134
    def exists?
135
      Dir.exist?(@db)
136
    end
137

138
    def status
139
      if exists?
140
        if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} status") == 0
141
          DatabaseStatus::RUNNING
142
        else
143
          DatabaseStatus::INACTIVE
144
        end
145
      else
146
        DatabaseStatus::NOT_FOUND
147
      end
148
    end
149

150
    def create_db_users(msf_pass, msftest_pass)
151
      puts 'Creating database users'
152
      run_psql("create user #{@options[:msf_db_user].shellescape} with password '#{msf_pass}'", @socket_directory)
153
      run_psql("create user #{@options[:msftest_db_user].shellescape} with password '#{msftest_pass}'", @socket_directory)
154
      run_psql("alter role #{@options[:msf_db_user].shellescape} createdb", @socket_directory)
155
      run_psql("alter role #{@options[:msftest_db_user].shellescape} createdb", @socket_directory)
156
      run_psql("alter role #{@options[:msf_db_user].shellescape} with password '#{msf_pass}'", @socket_directory)
157
      run_psql("alter role #{@options[:msftest_db_user].shellescape} with password '#{msftest_pass}'", @socket_directory)
158

159
      conn = PG.connect(host: @options[:db_host], dbname: 'postgres', port: @options[:db_port], user: @options[:msf_db_user], password: msf_pass)
160
      conn.exec("CREATE DATABASE #{@options[:msf_db_name]}")
161
      conn.exec("CREATE DATABASE #{@options[:msftest_db_name]}")
162
      conn.finish
163
    end
164

165
    def write_db_client_auth_config
166
      client_auth_config = "#{@db}/pg_hba.conf"
167
      super(client_auth_config)
168
    end
169

170
    def self.requirements
171
      %w[psql pg_ctl initdb createdb]
172
    end
173

174
    protected
175

176
    def retry_until_truthy(timeout:)
177
      start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
178
      ending_time = start_time + timeout
179
      retry_count = 0
180
      while Process.clock_gettime(Process::CLOCK_MONOTONIC, :second) < ending_time
181
        result = yield
182
        return result if result
183

184
        retry_count += 1
185
        remaining_time_budget = ending_time - Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
186
        break if remaining_time_budget <= 0
187

188
        delay = 2**retry_count
189
        if delay >= remaining_time_budget
190
          delay = remaining_time_budget
191
          puts("Final attempt. Sleeping for the remaining #{delay} seconds out of total timeout #{timeout}") if @options[:debug]
192
        else
193
          puts("Sleeping for #{delay} seconds before attempting again") if @options[:debug]
194
        end
195

196
        sleep delay
197
      end
198

199
      nil
200
    end
201
  end
202
end
203

204