Fix DB initialization with no database YAML

lib/metasploit/framework/data_service.rb

@@ -27,11 +27,15 @@ module DataService
   include LootDataService
 
   def name
-    raise 'DataLService#name is not implemented';
+    raise 'DataService#name is not implemented';
   end
 
   def active
-    raise 'DataLService#active is not implemented';
+    raise 'DataService#active is not implemented';
+  end
+
+  def is_local?
+    raise 'DataService#is_local? is not implemented';
   end
 
   #
@@ -41,11 +45,14 @@ module DataService
     attr_reader :id
     attr_reader :name
     attr_reader :active
+    attr_reader :is_local
 
-    def initialize (id, name, active)
+    def initialize (id, name, active, is_local)
       self.id = id
       self.name = name
       self.active = active
+      self.is_local = is_local
+
     end
 
     #######
@@ -55,6 +62,7 @@ module DataService
     attr_writer :id
     attr_writer :name
     attr_writer :active
+    attr_writer :is_local
 
   end
 end

lib/metasploit/framework/data_service/proxy/core.rb

@@ -28,13 +28,13 @@ class DataProxy
   #
   def error
     return @error if (@error)
-    return @current_data_service.error if @current_data_service
-    return 'none'
+    return @current_data_service.error if @current_data_service && !@current_data_service.error.nil?
+    return 'unknown'
   end
 
   def is_local?
-    if (@current_data_service)
-      return (@current_data_service.name == 'local_db_service')
+    if @current_data_service
+      return @current_data_service.is_local?
     end
 
     return false
@@ -44,7 +44,7 @@ class DataProxy
   # Determines if the data service is active
   #
   def active
-    if (@current_data_service)
+    if @current_data_service
       return @current_data_service.active
     end
 
@@ -67,11 +67,11 @@ class DataProxy
   #
   def set_data_service(data_service_id, online=false)
     data_service = @data_services[data_service_id.to_i]
-    if (data_service.nil?)
+    if data_service.nil?
       raise "Data service with id: #{data_service_id} does not exist"
     end
 
-    if (!online && !data_service.active)
+    if !online && !data_service.active
       raise "Data service not online: #{data_service.name}, not setting as active"
     end
 
@@ -86,7 +86,8 @@ class DataProxy
     @data_services.each_key {|key|
       name = @data_services[key].name
       active = !@current_data_service.nil? && name == @current_data_service.name
-      services_metadata << Metasploit::Framework::DataService::Metadata.new(key, name, active)
+      is_local = @data_services[key].is_local?
+      services_metadata << Metasploit::Framework::DataService::Metadata.new(key, name, active, is_local)
     }
 
     services_metadata

lib/metasploit/framework/data_service/remote/http/core.rb

@@ -170,6 +170,10 @@ class RemoteHTTPDataService
     "remote_data_service: (#{@endpoint})"
   end
 
+  def is_local?
+    false
+  end
+
   def set_header(key, value)
     @headers = Hash.new() if @headers.nil?
 

lib/msf/core/db_manager.rb

@@ -103,6 +103,10 @@ class Msf::DBManager
     'local_db_service'
   end
 
+  def is_local?
+    true
+  end
+
   #
   # Attributes
   #
@@ -153,9 +157,6 @@ class Msf::DBManager
   #
   def initialize_database_support
     begin
-      # Database drivers can reset our KCODE, do not let them
-      $KCODE = 'NONE' if RUBY_VERSION =~ /^1\.8\./
-
       add_rails_engine_migration_paths
 
       @usable = true
@@ -190,8 +191,11 @@ class Msf::DBManager
     else
       configuration_pathname = Metasploit::Framework::Database.configurations_pathname(path: opts['DatabaseYAML'])
 
-      unless configuration_pathname.nil?
+      if configuration_pathname.nil?
+        self.error = "No database YAML file"
+      else
         if configuration_pathname.readable?
+          # parse specified database YAML file
           dbinfo = YAML.load_file(configuration_pathname) || {}
           dbenv  = opts['DatabaseEnv'] || Rails.env
           db     = dbinfo[dbenv]

lib/msf/core/db_manager/connection.rb

@@ -64,9 +64,6 @@ module Msf::DBManager::Connection
       return false
     ensure
       after_establish_connection
-
-      # Database drivers can reset our KCODE, do not let them
-      $KCODE = 'NONE' if RUBY_VERSION =~ /^1\.8\./
     end
 
     true
@@ -139,9 +136,6 @@ module Msf::DBManager::Connection
     rescue ::Exception => e
       self.error = e
       elog("DB.disconnect threw an exception: #{e}")
-    ensure
-      # Database drivers can reset our KCODE, do not let them
-      $KCODE = 'NONE' if RUBY_VERSION =~ /^1\.8\./
     end
   end
 end

lib/msf/core/framework.rb

@@ -272,10 +272,12 @@ protected
   private
 
   def get_db
-    if !options['DisableDatabase']
+    unless options['DisableDatabase']
       db_manager = Msf::DBManager.new(self)
-      db_manager.init_db(options)
       options[:db_manager] = db_manager
+      unless options['SkipDatabaseInit']
+        db_manager.init_db(options)
+      end
     end
 
     Metasploit::Framework::DataService::DataProxy.new(options)

lib/msf/ui/console/command_dispatcher/db.rb

@@ -1653,8 +1653,8 @@ class Db
   # Database management
   #
   def db_check_driver
-    if(not framework.db.driver)
-      print_error("No database driver installed. Try 'gem install pg'")
+    unless framework.db.driver
+      print_error("No database driver installed.")
       return false
     end
     true
@@ -1909,6 +1909,12 @@ class Db
   #######
 
   def add_data_service(*args)
+    # database is required to use Mdm objects
+    unless framework.db.active
+      print_error("Database not connected; connect to an existing database with db_connect before using data_services")
+      return
+    end
+
     protocol = "http"
     port = 8080
     https_opts = {}

lib/msf/ui/console/driver.rb

@@ -4,6 +4,8 @@ require 'msf/base'
 require 'msf/ui'
 require 'msf/ui/console/framework_event_manager'
 require 'msf/ui/console/command_dispatcher'
+require 'msf/ui/console/command_dispatcher/db'
+require 'msf/ui/console/command_dispatcher/creds'
 require 'msf/ui/console/table'
 require 'find'
 require 'erb'
@@ -32,7 +34,9 @@ class Driver < Msf::Ui::Driver
   CommandDispatchers = [
     CommandDispatcher::Modules,
     CommandDispatcher::Jobs,
-    CommandDispatcher::Resource
+    CommandDispatcher::Resource,
+    CommandDispatcher::Db,
+    CommandDispatcher::Creds
   ]
 
   #
@@ -122,17 +126,13 @@ class Driver < Msf::Ui::Driver
       enstack_dispatcher(dispatcher)
     end
 
-    if framework.db && framework.db.active
-        require 'msf/ui/console/command_dispatcher/db'
-        enstack_dispatcher(CommandDispatcher::Db)
-        require 'msf/ui/console/command_dispatcher/creds'
-        enstack_dispatcher(CommandDispatcher::Creds)
-    else
+    if !framework.db || !framework.db.active
       print_error("***")
       if framework.db.error == "disabled"
         print_error("* WARNING: Database support has been disabled")
       else
-        print_error("* WARNING: No database support: #{framework.db.error.class} #{framework.db.error}")
+        error_msg = "#{framework.db.error.class.is_a?(String) ? "#{framework.db.error.class} " : nil}#{framework.db.error}"
+        print_error("* WARNING: No database support: #{error_msg}")
       end
       print_error("***")
     end
@@ -149,46 +149,6 @@ class Driver < Msf::Ui::Driver
     # Whether or not to confirm before exiting
     self.confirm_exit = opts['ConfirmExit']
 
-    # Parse any specified database.yml file
-    if framework.db.usable and not opts['SkipDatabaseInit']
-
-      # Append any migration paths necessary to bring the database online
-      if opts['DatabaseMigrationPaths']
-        opts['DatabaseMigrationPaths'].each do |migrations_path|
-          ActiveRecord::Migrator.migrations_paths << migrations_path
-        end
-      end
-
-      if framework.db.connection_established?
-        framework.db.after_establish_connection
-      else
-        configuration_pathname = Metasploit::Framework::Database.configurations_pathname(path: opts['DatabaseYAML'])
-
-        unless configuration_pathname.nil?
-          if configuration_pathname.readable?
-            dbinfo = YAML.load_file(configuration_pathname) || {}
-            dbenv  = opts['DatabaseEnv'] || Rails.env
-            db     = dbinfo[dbenv]
-          else
-            print_error("Warning, #{configuration_pathname} is not readable. Try running as root or chmod.")
-          end
-
-          if not db
-            print_error("No database definition for environment #{dbenv}")
-          else
-            framework.db.connect(db)
-          end
-        end
-      end
-
-      # framework.db.active will be true if after_establish_connection ran
-      # directly when connection_established? was already true or if
-      # framework.db.connect called after_establish_connection.
-      if !!framework.db.error
-        print_error("Failed to connect to the database: #{framework.db.error}")
-      end
-    end
-
     # Initialize the module paths only if we didn't get passed a Framework instance and 'DeferModuleLoads' is false
     unless opts['Framework'] || opts['DeferModuleLoads']
       # Configure the framework module paths