Land #10532, enhance db_connect for data services

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

@@ -78,6 +78,19 @@ class DataProxy
     end
   end
 
+  def delete_current_data_service
+    @data_services.each do |id, ds|
+      if ds == @current_data_service
+        if id == 1
+          raise "Unable to delete the local data service. Please use db_disconnect."
+        else
+          @data_services.delete(id)
+          @current_data_service = @data_services[1]
+        end
+      end
+    end
+  end
+
   #
   # Set the data service to be used
   #
@@ -185,6 +198,7 @@ class DataProxy
         @error = 'disabled'
       end
     rescue => e
+      @error = e
       raise "Unable to initialize data service: #{e.message}"
     end
   end
@@ -193,6 +207,13 @@ class DataProxy
     raise "Invalid data_service: #{data_service.class}, not of type Metasploit::Framework::DataService" unless data_service.is_a? (Metasploit::Framework::DataService)
     raise 'Cannot register null data service data_service' unless data_service
     raise 'Data Service already exists' if data_service_exist?(data_service)
+    # Raising an error for local DB causes startup to fail if there is a DB configured but we are unable to connect
+    # TODO: The check here shouldn't be dependent on if the data_service is local or not. We shouldn't
+    # connect to any data service if it is not online/active. This can likely be fixed by making a true
+    # LocalDataService instead of using DBManager.
+    unless data_service.is_local?
+      raise 'Data Service does not appear to be responding' unless data_service.active
+    end
   end
 
   def data_service_exist?(data_service)

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

@@ -39,7 +39,6 @@ module WorkspaceDataProxy
       else
         # This is mostly a failsafe to prevent bad things from happening. @current_workspace should always be set
         # outside of here, but this will save us from crashes/infinite loops if that happens
-        warn "@current_workspace was not set. Setting to default_workspace: #{default_workspace.name}"
         @current_workspace = default_workspace
       end
     rescue => e

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

@@ -22,6 +22,8 @@ class RemoteHTTPDataService
   DELETE_REQUEST = 'DELETE'
   PUT_REQUEST = 'PUT'
 
+  attr_reader :endpoint, :https_opts, :api_token
+
   #
   # @param [String] endpoint A valid http or https URL. Cannot be nil
   #
@@ -71,6 +73,10 @@ class RemoteHTTPDataService
     'none'
   end
 
+  def driver
+    'http'
+  end
+
   #
   # POST data to the HTTP endpoint and don't wait for the endpoint to process the data before getting a response
   #
@@ -289,7 +295,7 @@ class RemoteHTTPDataService
       if @endpoint.is_a?(URI::HTTPS)
         http.use_ssl = true
         http.verify_mode = OpenSSL::SSL::VERIFY_PEER
-        unless @https_opts.empty?
+        if @https_opts && !@https_opts.empty?
           if @https_opts[:skip_verify]
             http.verify_mode = OpenSSL::SSL::VERIFY_NONE
           else

lib/msf/base/config.rb

@@ -243,6 +243,14 @@ class Config < Hash
     self.new.save(opts)
   end
 
+  # Deletes the specified config group from the ini file
+  #
+  # @param group [String] The name of the group to remove
+  # @return [void]
+  def self.delete_group(group)
+    self.new.delete_group(group)
+  end
+
   # Updates the config class' self with the default hash.
   #
   # @return [Hash] the updated Hash.
@@ -424,6 +432,17 @@ class Config < Hash
     ini.to_file
   end
 
+  # Deletes the specified config group from the ini file
+  #
+  # @param group [String] The name of the group to remove
+  # @return [void]
+  def delete_group(group)
+    ini = Rex::Parser::Ini.new(config_file)
+
+    ini.delete(group)
+
+    ini.to_file
+  end
 end
 
 end

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

@@ -19,6 +19,8 @@ class Db
   include Msf::Ui::Console::CommandDispatcher
   include Msf::Ui::Console::CommandDispatcher::Common
 
+  DB_CONFIG_PATH = 'framework/database'
+
   #
   # The dispatcher's name.
   #
@@ -31,9 +33,11 @@ class Db
   #
   def commands
     base = {
-      "db_connect"    => "Connect to an existing database",
+      "db_connect"    => "Connect to an existing data service",
-      "db_disconnect" => "Disconnect from the current database instance",
+      "db_disconnect" => "Disconnect from the current data service",
-      "db_status"     => "Show the current database status",
+      "db_status"     => "Show the current data service status",
+      "db_save"       => "Save the current data service connection as the default to reconnect on startup",
+      "db_remove"     => "Remove the saved data service entry"
     }
 
     more = {
@@ -47,7 +51,6 @@ class Db
       "db_export"     => "Export a file containing the contents of the database",
       "db_nmap"       => "Executes nmap and records the output automatically",
       "db_rebuild_cache" => "Rebuilds the database-stored module cache",
-      "data_services" => "Command to add, list and set a data service",
     }
 
     # Always include commands that only make sense when connected.
@@ -82,28 +85,6 @@ class Db
     true
   end
 
-  def cmd_data_services(*args)
-    while (arg = args.shift)
-      case arg
-        when '-h', '--help'
-          data_service_help
-          return
-        when '-a', '--add'
-          add_data_service(*args)
-          return
-        when '-d', '--delete'
-          delete_data_service(args.shift)
-          return
-        when '-s', '--set'
-          set_data_service(args.shift)
-          return
-      end
-    end
-
-    list_data_services
-  end
-
-
   def cmd_workspace_help
     print_line "Usage:"
     print_line "    workspace                  List workspaces"
@@ -1203,7 +1184,7 @@ class Db
   end
 
   def cmd_loot_help
-    print_line "Usage: loot <options>"
+    print_line "Usage: loot [options]"
     print_line " Info: loot [-h] [addr1 addr2 ...] [-t <type1,type2>]"
     print_line "  Add: loot -f [fname] -i [info] -a [addr1 addr2 ...] -t [type]"
     print_line "  Del: loot -d [addr1 addr2 ...]"
@@ -1722,75 +1703,139 @@ class Db
     return if not db_check_driver
 
     if framework.db.connection_established?
-      cdb = ''
+      print_connection_info
-      ::ActiveRecord::Base.connection_pool.with_connection do |conn|
-        if conn.respond_to?(:current_database)
-          cdb = conn.current_database
-        end
-      end
-      print_status("#{framework.db.driver} connected to #{cdb}")
     else
       print_status("#{framework.db.driver} selected, no connection")
     end
   end
 
   def cmd_db_connect_help
-    # Help is specific to each driver
+    print_line("   Usage: db_connect <user:pass>@<host:port>/<database>")
-    cmd_db_connect("-h")
+    print_line("      OR: db_connect -y [path/to/database.yml]")
+    print_line("      OR: db_connect [options] <http|https>://<host:port>")
+    print_line("Examples:")
+    print_line("       db_connect user@metasploit3")
+    print_line("       db_connect user:pass@192.168.0.2/metasploit3")
+    print_line("       db_connect user:pass@192.168.0.2:1500/metasploit3")
+    print_line("       db_connect http://localhost:8080")
+    print_line("       db_connect -c ~/cert.pem -t 6a7a74c1a5003802c955ead1bbddd4ab1b05a7f2940b4732d34bfc555bc6e1c5d7611a497b29e8f0 https://localhost:8080")
+    print_line("       db_connect --name LA-server http://laoffice.org:8080")
+    print_line("       db_connect LA-server")
+    print_line(" ")
+    print_line("   OPTIONS:")
+    print_line("       -l,--list-services List the available data services that have been previously saved.")
+    print_line("       -y,--yaml          Connect to the data service specified in the provided database.yml file.")
+    print_line("       -n,--name          Name used to store the connection. Providing an existing name will overwrite the settings for that connection.")
+    print_line("       -c,--cert          Certificate file matching the remote data server's certificate. Needed when using self-signed SSL cert.")
+    print_line("       -t,--token         The API token used to authenticate to the remote data service.")
+    print_line("       --skip-verify      Skip validating authenticity of server's certificate (NOT RECOMMENDED).")
   end
 
   def cmd_db_connect(*args)
     return if not db_check_driver
-    if args[0] != '-h' && framework.db.connection_established?
+
-      cdb = ''
+    opts = {}
-      ::ActiveRecord::Base.connection_pool.with_connection do |conn|
+    https_opts = {}
-        if conn.respond_to?(:current_database)
+    while (arg = args.shift)
-          cdb = conn.current_database
+      case arg
+        when '-h', '--help'
+          cmd_db_connect_help
+          return
+        when '-y', '--yaml'
+          yaml_file = args.shift
+        when '-c', '--cert'
+          https_opts[:cert] = args.shift
+        when '-t', '--token'
+          opts[:api_token] = args.shift
+        when '-l', '--list-services'
+          list_saved_data_services
+          return
+      when '-n', '--name'
+          name = args.shift
+          if name =~ /\/|\[|\]/
+            print_error "Provided name contains an invalid character. Aborting connection."
+            return
+          end
+        when '--skip-verify'
+          https_opts[:skip_verify] = true
+      else
+        found_name = data_service_search(arg)
+        if found_name
+          opts = load_db_config(found_name)
+        else
+          opts[:url] = arg
         end
       end
-      print_error("#{framework.db.driver} already connected to #{cdb}")
+    end
-      print_error('Run db_disconnect first if you wish to connect to a different database')
+
+    opts[:https_opts] = https_opts unless https_opts.empty?
+
+    if !opts[:url] && !yaml_file
+      print_error 'A URL or saved data service name is required.'
+      print_line
+      cmd_db_connect_help
       return
     end
-    if (args[0] == "-y")
+
-      if (args[1] and not ::File.exist? ::File.expand_path(args[1]))
+    if opts[:url] =~ /http/
+      new_conn_type = 'http'
+    else
+      new_conn_type = framework.db.driver
+    end
+
+    # Currently only able to be connected to one DB at a time
+    if framework.db.connection_established?
+      # But the http connection still requires a local database to support AR, so we have to allow that
+      # Don't allow more than one HTTP service, though
+      if new_conn_type != 'http' || framework.db.get_services_metadata.count >= 2
+        print_error('Connection already established. Only one connection is allowed at a time.')
+        print_error('Run db_disconnect first if you wish to connect to a different data service.')
+        print_line
+        print_line 'Current connection information:'
+        print_connection_info
+        return
+      end
+    end
+
+    if yaml_file
+      if (yaml_file and not ::File.exist? ::File.expand_path(yaml_file))
         print_error("File not found")
         return
       end
-      file = args[1] || ::File.join(Msf::Config.get_config_root, "database.yml")
+      file = yaml_file || ::File.join(Msf::Config.get_config_root, "database.yml")
       file = ::File.expand_path(file)
       if (::File.exist? file)
         db = YAML.load(::File.read(file))['production']
         framework.db.connect(db)
-
+        print_line('Connected to the database specified in the YAML file.')
-        if framework.db.active and not framework.db.modules_cached
-          print_status("Rebuilding the module cache in the background...")
-          framework.threads.spawn("ModuleCacheRebuild", true) do
-            framework.db.update_all_module_details
-          end
-        end
-
         return
       end
     end
-    meth = "db_connect_#{framework.db.driver}"
+
+    meth = "db_connect_#{new_conn_type}"
     if(self.respond_to?(meth, true))
-      self.send(meth, *args)
+      self.send(meth, opts)
-      if framework.db.active and not framework.db.modules_cached
+    else
-        print_status("Rebuilding the module cache in the background...")
+      print_error("This database driver #{new_conn_type} is not currently supported")
-        framework.threads.spawn("ModuleCacheRebuild", true) do
+    end
-          framework.db.update_all_module_details
+
+    if framework.db.active
+      if !name || name.empty?
+        if found_name
+          name = found_name
+        else
+          name = Rex::Text.rand_text_alphanumeric(8)
         end
       end
-    else
+      save_db_to_config(framework.db, name)
-      print_error("This database driver #{framework.db.driver} is not currently supported")
+      @current_data_service = name
     end
   end
 
   def cmd_db_disconnect_help
     print_line "Usage: db_disconnect"
     print_line
-    print_line "Disconnect from the database."
+    print_line "Disconnect from the data service."
     print_line
   end
 
@@ -1802,8 +1847,25 @@ class Db
       return
     end
 
-    if (framework.db)
+    db_name = framework.db.name
-      framework.db.disconnect()
+
+    if framework.db.active
+      if framework.db.driver == 'http'
+        begin
+          framework.db.delete_current_data_service
+          local_db_url = build_postgres_url
+          local_name = data_service_search(local_db_url)
+          @current_data_service = local_name
+        rescue => e
+          print_error "Unable to disconnect from the data service: #{e.message}"
+        end
+      else
+        framework.db.disconnect
+        @current_data_service = nil
+      end
+      print_line "Successfully disconnected from the data service: #{db_name}."
+    else
+      print_error "Not currently connected to a data service."
     end
   end
 
@@ -1827,6 +1889,101 @@ class Db
     print_line
   end
 
+  def cmd_db_save_help
+    print_line "Usage: db_save"
+    print_line
+    print_line "Save the current data service connection as the default to reconnect on startup."
+    print_line
+  end
+
+  def cmd_db_save(*args)
+    while (arg = args.shift)
+      case arg
+        when '-h', '--help'
+          cmd_db_save_help
+          return
+      end
+    end
+
+    if !framework.db.active || !@current_data_service
+      print_error "Not currently connected to a data service that can be saved."
+      return
+    end
+
+    begin
+      Msf::Config.save(DB_CONFIG_PATH => { 'default_db' => @current_data_service })
+      print_line "Successfully saved data service as default: #{@current_data_service}"
+    rescue ArgumentError => e
+      print_error e.message
+    end
+  end
+
+  def save_db_to_config(database, database_name)
+    if database_name =~ /\/|\[|\]/
+      raise ArgumentError, 'Data service name contains an invalid character.'
+    end
+    config_path = "#{DB_CONFIG_PATH}/#{database_name}"
+    config_opts = {}
+    if !database.is_local?
+      begin
+        config_opts['url'] = database.endpoint
+        if database.https_opts
+          config_opts['cert'] = database.https_opts[:cert] if database.https_opts[:cert]
+          config_opts['skip_verify'] = true if database.https_opts[:skip_verify]
+        end
+        if database.api_token
+          config_opts['api_token'] = database.api_token
+        end
+        Msf::Config.save(config_path => config_opts)
+      rescue => e
+        print_error "There was an error saving the data service configuration: #{e.message}"
+      end
+    else
+      url = build_postgres_url
+      config_opts['url'] = url
+      Msf::Config.save(config_path => config_opts)
+    end
+  end
+
+  def cmd_db_remove_help
+    print_line "Usage: db_remove <name>"
+    print_line
+    print_line "Delete the specified saved data service."
+    print_line
+  end
+
+  def cmd_db_remove(*args)
+    if args[0] == '-h' || args[0] == '--help' || args[0].nil? || args[0].empty?
+      cmd_db_remove_help
+      return
+    end
+    delete_db_from_config(args[0])
+  end
+
+  def delete_db_from_config(db_name)
+    conf = Msf::Config.load
+    db_path = "#{DB_CONFIG_PATH}/#{db_name}"
+    if conf[db_path]
+      clear_default_db if conf[DB_CONFIG_PATH]['default_db'] && conf[DB_CONFIG_PATH]['default_db'] == db_name
+      Msf::Config.delete_group(db_path)
+      print_line "Successfully deleted data service: #{db_name}"
+    else
+      print_line "Unable to locate saved data service with name #{db_name}."
+    end
+  end
+
+  def clear_default_db
+    conf = Msf::Config.load
+    if conf[DB_CONFIG_PATH] && conf[DB_CONFIG_PATH]['default_db']
+      updated_opts = conf[DB_CONFIG_PATH]
+      updated_opts.delete('default_db')
+      Msf::Config.save(DB_CONFIG_PATH => updated_opts)
+      print_line "Cleared the default data service."
+    else
+      print_line "No default data service was configured."
+    end
+  end
+
   def db_find_tools(tools)
     missed  = []
     tools.each do |name|
@@ -1848,18 +2005,8 @@ class Db
   #
   # Connect to an existing Postgres database
   #
-  def db_connect_postgresql(*args)
+  def db_connect_postgresql(cli_opts)
-    if(args[0] == nil or args[0] == "-h" or args[0] == "--help")
+    info = db_parse_db_uri_postgresql(cli_opts[:url])
-      print_status("   Usage: db_connect <user:pass>@<host:port>/<database>")
-      print_status("      OR: db_connect -y [path/to/database.yml]")
-      print_status("Examples:")
-      print_status("       db_connect user@metasploit3")
-      print_status("       db_connect user:pass@192.168.0.2/metasploit3")
-      print_status("       db_connect user:pass@192.168.0.2:1500/metasploit3")
-      return
-    end
-
-    info = db_parse_db_uri_postgresql(args[0])
     opts = { 'adapter' => 'postgresql' }
 
     opts['username'] = info[:user] if (info[:user])
@@ -1896,8 +2043,29 @@ class Db
       opts['host'] = '127.0.0.1'
     end
 
-    if (not framework.db.connect(opts))
+    if framework.db.connect(opts) && framework.db.connection_established?
-      raise RuntimeError.new("Failed to connect to the database: #{framework.db.error}")
+      print_line "Connected to Postgres data service: #{info[:host]}/#{info[:name]}"
+    else
+      raise RuntimeError.new("Failed to connect to the Postgres data service: #{framework.db.error}")
+    end
+  end
+
+  def db_connect_http(opts)
+    # local database is required to use Mdm objects
+    unless framework.db.active
+      print_error("No local database connected. Please connect to a local database before connecting to a remote data service.")
+      return
+    end
+
+    uri = db_parse_db_uri_http(opts[:url])
+
+    remote_data_service = Metasploit::Framework::DataService::RemoteHTTPDataService.new(uri.to_s, opts)
+    begin
+      framework.db.register_data_service(remote_data_service)
+      print_line "Connected to HTTP data service: #{remote_data_service.name}"
+      framework.db.workspace = framework.db.default_workspace
+    rescue => e
+      raise RuntimeError.new("Failed to connect to the HTTP data service: #{e.message}")
     end
   end
 
@@ -1917,6 +2085,10 @@ class Db
     res
   end
 
+  def db_parse_db_uri_http(path)
+    URI.parse(path)
+  end
+
   #
   # Miscellaneous option helpers
   #
@@ -1963,97 +2135,89 @@ class Db
   private
   #######
 
-  def add_data_service(*args)
+  def print_connection_info
-    # database is required to use Mdm objects
+    cdb = ''
-    unless framework.db.active
+    if framework.db.driver == 'http'
-      print_error("Database not connected; connect to an existing database with db_connect before using data_services")
+      cdb = framework.db.name
-      return
+    else
-    end
+      ::ActiveRecord::Base.connection_pool.with_connection do |conn|
-
+        if conn.respond_to?(:current_database)
-    protocol = "http"
+          cdb = conn.current_database
-    port = 8080
+        end
-    opts = {}
-    https_opts = {}
-    while (arg = args.shift)
-      case arg
-        when '-p', '--port'
-          port = args.shift
-        when '-t', '--token'
-          opts[:api_token] = args.shift
-        when '-s', '--ssl'
-          protocol = "https"
-        when '-c', '--cert'
-          https_opts[:cert] = args.shift
-        when '--skip-verify'
-          https_opts[:skip_verify] = true
-        else
-          host = arg
       end
     end
+    output = "Connected to #{cdb}. Connection type: #{framework.db.driver}."
+    output += " Connection name: #{@current_data_service}." if @current_data_service
+    print_status(output)
+  end
 
-    if host.nil? || port.nil?
+  def data_service_search(search_criteria)
-      print_error("Host and port are required")
+    conf = Msf::Config.load
+    rv = nil
+
+    conf.each_pair do |k,v|
+      name = k.split('/').last
+      rv = name if name == search_criteria
+      rv = name if v.values.include?(search_criteria)
+    end
+    rv
+  end
+
+  def load_db_config(db_name)
+    conf = Msf::Config.load
+    conf_options = conf["#{DB_CONFIG_PATH}/#{db_name}"]
+    opts = {}
+    https_opts = {}
+    if conf_options
+      opts[:url] = conf_options['url'] if conf_options['url']
+      opts[:api_token] = conf_options['api_token'] if conf_options['api_token']
+      https_opts[:cert] = conf_options['cert'] if conf_options['cert']
+      https_opts[:skip_verify] = conf_options['skip_verify'] if conf_options['skip_verify']
+    else
+      print_error "Unable to locate saved data service with name '#{db_name}'"
       return
     end
 
     opts[:https_opts] = https_opts unless https_opts.empty?
-    endpoint = "#{protocol}://#{host}:#{port}"
+    opts
-    remote_data_service = Metasploit::Framework::DataService::RemoteHTTPDataService.new(endpoint, opts)
-    begin
-      framework.db.register_data_service(remote_data_service)
-      print_line "Registered data service: #{remote_data_service.name}"
-      framework.db.workspace = framework.db.default_workspace
-    rescue => e
-      print_error "There was a problem registering the remote data service: #{e.message}"
-    end
   end
 
-  def delete_data_service(service_id)
+  def list_saved_data_services
-    begin
+    conf = Msf::Config.load
-      data_service = framework.db.delete_data_service(service_id)
+    default = nil
-      framework.db.workspace = framework.db.default_workspace
+    tbl = Rex::Text::Table.new({
-      data_service
+                                   'Header'    => 'Data Services',
-    rescue => e
+                                   'Columns'   => ['current', 'name', 'url', 'default?'],
-      print_error "Unable to delete data service: #{e.message}"
+                                   'SortIndex' => 1
-    end
+                               })
-  end
 
-  def set_data_service(service_id)
+    conf.each_pair do |k,v|
-    begin
+      if k =~ /#{DB_CONFIG_PATH}/
-      data_service = framework.db.set_data_service(service_id)
+        default = v['default_db'] if v['default_db']
-      framework.db.workspace = framework.db.default_workspace
+        name = k.split('/').last
-      data_service
+        next if name == 'database' # Data service information is not stored in 'framework/database', just metadata
-    rescue => e
+        url = v['url']
-      print_error "Unable to set data service: #{e.message}"
+        current = ''
-    end
+        current = '*' if name == @current_data_service
-  end
+        default_output = ''
-
+        default_output = '*' if name == default
-  def list_data_services()
+        line = [current, name, url, default_output]
-    framework.db.get_services_metadata.each {|metadata|
+        tbl << line
-      out = "id: #{metadata.id}, name: #{metadata.name}"
-      if metadata.active
-        out += " [active]"
       end
-      print_line out
+    end
-    }
+    print_line
+    print_line tbl.to_s
   end
 
-  def data_service_help
+  def build_postgres_url
-    print_line "Usage: data_services [ options ] - list data services by default"
+    conn_params = ActiveRecord::Base.connection_config
-    print_line
+    url = ""
-    print_line "OPTIONS:"
+    url += "#{conn_params[:username]}" if conn_params[:username]
-
+    url += ":#{conn_params[:password]}" if conn_params[:password]
-    print_line "  -h, --help                    Show this help information."
+    url += "@#{conn_params[:host]}" if conn_params[:host]
-    print_line "  -d, --delete <id>             Delete the data service by identifier."
+    url += ":#{conn_params[:port]}" if conn_params[:port]
-    print_line "  -s, --set <id>                Set the active data service by identifier."
+    url += "/#{conn_params[:database]}" if conn_params[:database]
-    print_line "  -a, --add [ options ] <host>  Add a new data service"
+    url
-    print_line "  Add Data Service Options:"
-    print_line "  -p, --port <port>   The port the data service is listening on. Default is 8080."
-    print_line "  -t, --token <token> API Token for MSF web service"
-    print_line "  -s, --ssl           Enable SSL. Required for HTTPS data services."
-    print_line "  -c, --cert          Certificate file matching the server's certificate. Needed when using self-signed SSL cert."
-    print_line "  --skip-verify       Skip validating authenticity of server's certificate. NOT RECOMMENDED."
-    print_line
   end
 
   def print_msgs(status_msg, error_msg)

lib/msf/ui/console/driver.rb

@@ -24,6 +24,7 @@ class Driver < Msf::Ui::Driver
 
   ConfigCore  = "framework/core"
   ConfigGroup = "framework/ui/console"
+  DbConfigGroup = "framework/database"
 
   DefaultPrompt     = "%undmsf5%clr"
   DefaultPromptChar = "%clr>"
@@ -129,6 +130,8 @@ class Driver < Msf::Ui::Driver
       enstack_dispatcher(dispatcher)
     end
 
+    load_db_config(opts['Config'])
+
     if !framework.db || !framework.db.active
       print_error("***")
       if framework.db.error == "disabled"
@@ -224,6 +227,38 @@ class Driver < Msf::Ui::Driver
     end
   end
 
+  def load_db_config(path=nil)
+    begin
+      conf = Msf::Config.load(path)
+    rescue
+      wlog("Failed to load configuration: #{$!}")
+      return
+    end
+
+    if conf.group?(DbConfigGroup)
+      conf[DbConfigGroup].each_pair do |k, v|
+        if k.downcase == 'default_db'
+          ilog "Default data service found. Attempting to connect..."
+          default_db_config_path = "#{DbConfigGroup}/#{v}"
+          default_db = conf[default_db_config_path]
+          if default_db
+            connect_string = "db_connect #{v}"
+
+            if framework.db.active && default_db['url'] !~ /http/
+              ilog "Existing local data connection found. Disconnecting first."
+              run_single("db_disconnect")
+            end
+
+            run_single(connect_string)
+          else
+            elog "Config entry for '#{default_db_config_path}' could not be found. Config file might be corrupt."
+            return
+          end
+        end
+      end
+    end
+  end
+
   #
   # Loads configuration for the console.
   #
@@ -333,7 +368,10 @@ class Driver < Msf::Ui::Driver
         print_warning("\t#{path}: #{error}")
       end
     end
-    framework.db.workspace = framework.db.default_workspace if framework.db && framework.db.active
+
+    if framework.db && framework.db.active
+      framework.db.workspace = framework.db.default_workspace unless framework.db.workspace
+    end
 
     framework.events.on_ui_start(Msf::Framework::Revision)
 

spec/lib/msf/ui/console/command_dispatcher/db_spec.rb

@@ -147,7 +147,7 @@ RSpec.describe Msf::Ui::Console::CommandDispatcher::Db do
       it "should show a help message" do
         db.cmd_loot "-h"
         expect(@output).to match_array [
-          "Usage: loot <options>",
+          "Usage: loot [options]",
           " Info: loot [-h] [addr1 addr2 ...] [-t <type1,type2>]",
           "  Add: loot -f [fname] -i [info] -a [addr1 addr2 ...] -t [type]",
           "  Del: loot -d [addr1 addr2 ...]",