Ensure msfdb reinit can be used for starting the database

lib/msfdb_helpers/db_interface.rb

@@ -13,10 +13,6 @@ def delete
       raise NotImplementedError
     end
 
-    def reinit
-      raise NotImplementedError
-    end
-
     def start
       raise NotImplementedError
     end

lib/msfdb_helpers/pg_ctl.rb

@@ -12,21 +12,6 @@ def initialize(db_path:, options:, localconf:, db_conf:)
     end
 
     def init(msf_pass, msftest_pass)
-      if Dir.exist?(@db)
-        puts "Found a database at #{@db}, checking to see if it is started"
-        start
-        return
-      end
-
-      if File.exist?(@db_conf) && !@options[:delete_existing_data]
-        if !load_db_config
-          puts 'Failed to load existing database config. Please reinit and overwrite the file.'
-          return
-        end
-      else
-        write_db_config
-      end
-
       puts "Creating database at #{@db}"
       Dir.mkdir(@db)
       run_cmd("initdb --auth-host=trust --auth-local=trust -E UTF8 #{@db.shellescape}")
@@ -44,7 +29,7 @@ def init(msf_pass, msftest_pass)
     end
 
     def delete
-      if Dir.exist?(@db)
+      if exists?
         stop
 
         if @options[:delete_existing_data]
@@ -53,20 +38,15 @@ def delete
         end
 
         if @options[:delete_existing_data]
-          File.delete(@db_conf)
+          FileUtils.rm_r(@db_conf, force: true)
         end
       else
         puts "No data at #{@db}, doing nothing"
       end
     end
 
-    def reinit(msf_pass, msftest_pass)
-      delete
-      init(msf_pass, msftest_pass)
-    end
-
     def start
-      if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} status") == 0
+      if status == DatabaseStatus::RUNNING
         puts "Database already started at #{@db}"
         return true
       end
@@ -84,7 +64,7 @@ def start
     end
 
     def stop
-      if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} status") == 0
+      if status == DatabaseStatus::RUNNING
         puts "Stopping database at #{@db}"
         run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} stop")
       else
@@ -97,15 +77,19 @@ def restart
       start
     end
 
+    def exists?
+      Dir.exist?(@db)
+    end
+
     def status
-      if Dir.exist?(@db)
+      if exists?
         if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} status") == 0
-          puts "Database started at #{@db}"
+          DatabaseStatus::RUNNING
         else
-          puts "Database is not running at #{@db}"
+          DatabaseStatus::INACTIVE
         end
       else
-        puts "No database found at #{@db}"
+        DatabaseStatus::NOT_FOUND
       end
     end
 

lib/msfdb_helpers/pg_ctlcluster.rb

@@ -15,21 +15,6 @@ def initialize(db_path:, options:, localconf:, db_conf:)
     end
 
     def init(msf_pass, msftest_pass)
-      if Dir.exist?(@db)
-        puts "Found a database at #{@db}, checking to see if it is started"
-        start
-        return
-      end
-
-      if File.exist?(@db_conf) && !@options[:delete_existing_data]
-        if !load_db_config
-          puts 'Failed to load existing database config. Please reinit and overwrite the file.'
-          return
-        end
-      else
-        write_db_config
-      end
-
       puts "Creating database at #{@db}"
       Dir.mkdir(@db)
       FileUtils.mkdir_p(@pg_cluster_conf_root)
@@ -47,26 +32,21 @@ def init(msf_pass, msftest_pass)
     end
 
     def delete
-      if Dir.exist?(@db)
+      if exists?
         stop
 
         if @options[:delete_existing_data]
           puts "Deleting all data at #{@db}"
           run_cmd("pg_dropcluster #{@pg_version} #{@options[:msf_db_name].shellescape}")
           FileUtils.rm_rf(@db)
           FileUtils.rm_rf("#{@localconf}/.local/etc/postgresql")
-          File.delete(@db_conf)
+          FileUtils.rm_r(@db_conf, force: true)
         end
       else
         puts "No data at #{@db}, doing nothing"
       end
     end
 
-    def reinit(msf_pass, msftest_pass)
-      delete
-      init(msf_pass, msftest_pass)
-    end
-
     def start
       print "Starting database at #{@db}..."
       status = run_cmd("pg_ctlcluster #{@pg_version} #{@options[:msf_db_name].shellescape} start -- -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} -l #{@db.shellescape}/log")
@@ -91,15 +71,19 @@ def restart
       run_cmd("pg_ctlcluster #{@pg_version} #{@options[:msf_db_name].shellescape} reload -- -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape} -l #{@db.shellescape}/log")
     end
 
+    def exists?
+      Dir.exist?(@db)
+    end
+
     def status
-      if Dir.exist?(@db)
+      if exists?
         if run_cmd("pg_ctlcluster #{@pg_version} #{@options[:msf_db_name].shellescape} status -- -o \"-p #{@options[:db_port]}\" -D #{@db.shellescape}") == 0
-          puts "Database started at #{@db}"
+          DatabaseStatus::RUNNING
         else
-          puts "Database is not running at #{@db}"
+          DatabaseStatus::INACTIVE
         end
       else
-        puts "No database found at #{@db}"
+        DatabaseStatus::NOT_FOUND
       end
     end
 
@@ -137,6 +121,5 @@ def create_db_users(msf_pass, msftest_pass)
       conn.exec("CREATE DATABASE #{@options[:msftest_db_name]}")
       conn.finish
     end
-
   end
 end

lib/msfdb_helpers/standalone.rb

@@ -19,46 +19,42 @@ def initialize(options:, db_conf:, connection_string:)
     end
 
     def init(msf_pass, msftest_pass)
-      @conn.exec("create user #{@options[:msf_db_user]} with password '#{msf_pass}'")
-      @conn.exec("create user #{@options[:msftest_db_user]} with password '#{msftest_pass}'")
-      @conn.exec("alter role #{@options[:msf_db_user]} createdb")
-      @conn.exec("alter role #{@options[:msftest_db_user]} createdb")
-      @conn.exec("alter role #{@options[:msf_db_user]} with password '#{msf_pass}'")
-      @conn.exec("alter role #{@options[:msftest_db_user]} with password '#{msftest_pass}'")
-      @conn.exec("CREATE DATABASE #{@options[:msf_db_name]}")
-      @conn.exec("CREATE DATABASE #{@options[:msftest_db_name]}")
-      @conn.finish
+      create_db_users(msf_pass, msftest_pass)
     end
 
     def delete
       @conn.exec("DROP DATABASE IF EXISTS #{@options[:msf_db_name]};")
       @conn.exec("DROP DATABASE IF EXISTS #{@options[:msftest_db_name]};")
       @conn.exec("DROP USER IF EXISTS #{@options[:msf_db_user]};")
       @conn.exec("DROP USER IF EXISTS #{@options[:msftest_db_user]};")
-      if File.exist?(@db_conf)
-        File.delete(@db_conf)
-      end
-    end
-
-    def reinit(msf_pass, msftest_pass)
-      delete
-      init(msf_pass, msftest_pass)
+      FileUtils.rm_r(@db_conf, force: true)
     end
 
     def start
-      raise NotImplementedError
+      true
     end
 
     def stop
-      raise NotImplementedError
+      puts 'A standalone database cannot be stopped by msfdb'
+      false
     end
 
     def restart
       raise NotImplementedError
     end
 
+    def exists?
+      !@conn.nil?
+    end
+
     def status
-      raise NotImplementedError
+      # Search for the database name
+      is_initialized = @conn.exec_params('select * from pg_catalog.pg_database where datname = $1', [@options[:msf_db_name]]).any?
+      if !is_initialized
+        DatabaseStatus::NEEDS_INIT
+      else
+        DatabaseStatus::RUNNING
+      end
     end
 
     def write_db_client_auth_config
@@ -69,5 +65,18 @@ def self.requirements
       []
     end
 
+    private
+
+    def create_db_users(msf_pass, msftest_pass)
+      @conn.exec("create user #{@options[:msf_db_user]} with password '#{msf_pass}'")
+      @conn.exec("create user #{@options[:msftest_db_user]} with password '#{msftest_pass}'")
+      @conn.exec("alter role #{@options[:msf_db_user]} createdb")
+      @conn.exec("alter role #{@options[:msftest_db_user]} createdb")
+      @conn.exec("alter role #{@options[:msf_db_user]} with password '#{msf_pass}'")
+      @conn.exec("alter role #{@options[:msftest_db_user]} with password '#{msftest_pass}'")
+      @conn.exec("CREATE DATABASE #{@options[:msf_db_name]}")
+      @conn.exec("CREATE DATABASE #{@options[:msftest_db_name]}")
+      @conn.finish
+    end
   end
 end

msfdb

@@ -131,29 +131,39 @@ end
 
 def status_db
   update_db_port
-  @db_driver.status
-end
 
-def started_db
-  @db_driver.start
+  case @db_driver.status
+  when DatabaseStatus::RUNNING
+    puts "Database started"
+  when DatabaseStatus::INACTIVE
+    puts "Database found, but is not running"
+  when DatabaseStatus::NEEDS_INIT
+    puts "Database found, but needs initialized"
+  when DatabaseStatus::NOT_FOUND
+    puts "No database found"
+  end
 end
 
 def start_db
-  if !Dir.exist?(@db)
-    print_error "No database found at #{@db}."
-    print_error "Has the database been initialized with \"msfdb init\" or \"msfdb init --component database\"?"
+  case @db_driver.status
+  when DatabaseStatus::NOT_FOUND
+    print_error 'No database found.'
+    return
+  when DatabaseStatus::NEEDS_INIT
+    print_error 'Has the database been initialized with "msfdb init" or "msfdb init --component database"?'
     return
   end
 
   update_db_port
+  db_started = @db_driver.start
 
-  if !started_db
+  if !db_started
     last_log = tail("#{@db}/log")
     puts last_log
     if last_log =~ /not compatible/
       puts 'Please attempt to upgrade the database manually using pg_upgrade.'
     end
-    print_error "Your database may be corrupt. Try reinitializing."
+    print_error 'Your database may be corrupt. Try reinitializing.'
   end
 end
 
@@ -167,14 +177,30 @@ def restart_db
 end
 
 def init_db
+  case @db_driver.status
+  when DatabaseStatus::RUNNING
+    puts 'Existing database running'
+    return
+  when DatabaseStatus::INACTIVE
+    puts 'Existing database found, attempting to start it'
+    @db_driver.start
+    return
+  end
+
+  if @db_driver.exists? && !@options[:delete_existing_data]
+    if !load_db_config
+      puts 'Failed to load existing database config. Please reinit and overwrite the file.'
+      return
+    end
+  end
+
   # Generate new database passwords if not already assigned
   @msf_pass ||= pw_gen
   @msftest_pass ||= pw_gen
 
   @db_driver.init(@msf_pass, @msftest_pass)
   write_db_config
 
-
   puts 'Creating initial database schema'
   Dir.chdir(@framework) do
     @db_driver.run_cmd('bundle exec rake db:migrate')
@@ -314,6 +340,13 @@ class WebServicePIDStatus
   NO_PID_FILE = 2
 end
 
+class DatabaseStatus
+  RUNNING = 0
+  INACTIVE = 1
+  NOT_FOUND = 2
+  NEEDS_INIT = 3
+end
+
 def web_service_pid
   File.file?(@ws_pid) ? tail(@ws_pid) : nil
 end
@@ -626,14 +659,14 @@ def run_msfconsole_command(cmd)
 end
 
 def persist_data_service
-  puts 'Persisting data service credentials in msfconsole'
+  puts 'Persisting http web data service credentials in msfconsole'
   # execute msfconsole commands to add and persist the data service connection
   cmd = "./msfconsole -qx '#{get_db_connect_command}; db_save; exit'"
   run_msfconsole_command(cmd)
 end
 
 def clear_default_data_service
-  puts 'Clearing data service credentials in msfconsole'
+  puts 'Clearing http web data service credentials in msfconsole'
   # execute msfconsole commands to clear the default data service connection
   cmd = "./msfconsole -qx 'db_disconnect --clear; exit'"
   run_msfconsole_command(cmd)
@@ -943,6 +976,10 @@ def prompt_for_component(command)
     return :all
   end
 
+  if command == :stop && web_service_pid_status != WebServicePIDStatus::RUNNING
+    return :database
+  end
+
   enable_webservice = ask_yn("Would you like to #{command} the webservice? (Not Required)", default: 'no')
   if enable_webservice
     :all