Merge pull request #7 from clee-r7/MS-2891

Land #7, fix test issues

Author: mkienow-r7

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

@@ -1,8 +1,6 @@
-require 'singleton'
 require 'open3'
 require 'rex/ui'
 require 'rex/logging'
-require 'msf/core/db_manager'
 require 'metasploit/framework/data_service/remote/http/core'
 require 'metasploit/framework/data_service/remote/http/remote_service_endpoint'
 require 'metasploit/framework/data_service/proxy/data_proxy_auto_loader'
@@ -15,11 +13,17 @@ module Metasploit
 module Framework
 module DataService
 class DataProxy
-  include Singleton
   include DataProxyAutoLoader
 
   attr_reader :usable
 
+  def initialize(opts = {})
+    @data_services = {}
+    @data_service_id = 0
+    @usable = false
+    setup(opts)
+  end
+
   #
   # Returns current error state
   #
@@ -48,34 +52,6 @@ def active
     return false
   end
 
-  #
-  # Initializes the data service to be used - primarily on startup
-  #
-  def init(framework, opts)
-    @mutex.synchronize {
-      if (@initialized)
-        return
-      end
-
-      begin
-        if (opts['DisableDatabase'])
-          @error = 'disabled'
-          return
-        elsif (opts['DatabaseRemoteProcess'])
-          run_remote_db_process(opts)
-        else
-          run_local_db_process(framework, opts)
-        end
-        @usable = true
-        @initialized = true
-      rescue Exception => e
-        puts "Unable to initialize a dataservice #{e.message}"
-        return
-      end
-    }
-
-  end
-
   #
   # Registers a data service with the proxy and immediately
   # set as primary if online
@@ -131,6 +107,14 @@ def method_missing(method, *args, &block)
     end
   end
 
+  def respond_to?(method_name, include_private=false)
+    unless @data_service.nil?
+      return @data_service.respond_to?(method_name, include_private)
+    end
+
+    false
+  end
+
   #
   # Attempt to shutdown the local db process if it exists
   #
@@ -145,10 +129,6 @@ def exit_called
     end
   end
 
-  #########
-  protected
-  #########
-
   def get_data_service
     raise 'No registered data_service' unless @data_service
     return @data_service
@@ -158,12 +138,21 @@ def get_data_service
   private
   #######
 
-  def initialize
-    @data_services = {}
-    @data_service_id = 0
-    @usable = false
-    @initialized = false
-    @mutex = Mutex.new()
+  def setup(opts)
+    begin
+      db_manager = opts.delete(:db_manager)
+      if !db_manager.nil?
+        register_data_service(db_manager, true)
+        @usable = true
+      elsif opts['DatabaseRemoteProcess']
+        run_remote_db_process(opts)
+        @usable = true
+      else
+        @error = 'disabled'
+      end
+    rescue Exception => e
+      puts "Unable to initialize a dataservice #{e.message}"
+    end
   end
 
   def validate(data_service)
@@ -183,15 +172,6 @@ def data_service_exist?(data_service)
   end
 
 
-  def run_local_db_process(framework, opts)
-    puts 'Initializing local db process'
-    db_manager = Msf::DBManager.new(framework)
-    if (db_manager.usable and not opts['SkipDatabaseInit'])
-        register_data_service(db_manager, true)
-        db_manager.init_db(opts)
-    end
-  end
-
   def run_remote_db_process(opts)
     # started with no signal to prevent ctrl-c from taking out db
     db_script = File.join( Msf::Config.install_root, "msfdb -ns")

lib/msf/core/db_manager/connection.rb

@@ -18,7 +18,7 @@ def after_establish_connection
       migrate
 
       # Set the default workspace
-      framework.db.workspace = framework.db.default_workspace
+      self.workspace = self.default_workspace
     rescue ::Exception => exception
       self.error = exception
       elog("DB.connect threw an exception: #{exception}")

lib/msf/core/db_manager/host.rb

@@ -56,7 +56,7 @@ def each_host(wspace=workspace, &block)
 
   # Exactly like report_host but waits for the database to create a host and returns it.
   def find_or_create_host(opts)
-    host = get_host(opts)
+    host = get_host(opts.clone)
     return host unless host.nil?
 
     report_host(opts)

lib/msf/core/db_manager/workspace.rb

@@ -58,20 +58,23 @@ def workspace_associations_counts()
   end
 
   def delete_all_workspaces()
-    delete_workspaces(workspaces.map(&:name))
+    return delete_workspaces(workspaces.map(&:name))
   end
 
   def delete_workspaces(names)
+    status_msg = []
+    error_msg = []
+
     switched = false
     # Delete workspaces
     names.each do |name|
       workspace = framework.db.find_workspace(name)
       if workspace.nil?
-        print_error("Workspace not found: #{name}")
+        error << "Workspace not found: #{name}"
       elsif workspace.default?
         workspace.destroy
         workspace = framework.db.add_workspace(name)
-        print_status("Deleted and recreated the default workspace")
+        status_msg << 'Deleted and recreated the default workspace'
       else
         # switch to the default workspace if we're about to delete the current one
         if framework.db.workspace.name == workspace.name
@@ -80,10 +83,11 @@ def delete_workspaces(names)
         end
         # now destroy the named workspace
         workspace.destroy
-        print_status("Deleted workspace: #{name}")
+        status_msg << "Deleted workspace: #{name}"
       end
     end
-    print_status("Switched workspace: #{framework.db.workspace.name}") if switched
+    (status_msg << "Switched workspace: #{framework.db.workspace.name}") if switched
+    return status_msg, error_msg
   end
 
   #

lib/msf/core/framework.rb

@@ -199,7 +199,7 @@ def version
   # @return [Metasploit::Framework::DataService::DataProxy]
   def db
     synchronize {
-      @db ||= Metasploit::Framework::DataService::DataProxy.instance
+      @db ||= get_db
     }
   end
 
@@ -268,6 +268,19 @@ def search(match, logger: nil)
   attr_writer   :db # :nodoc:
   attr_writer   :uuid_db # :nodoc:
   attr_writer   :browser_profiles # :nodoc:
+
+  private
+
+  def get_db
+    if !options['DatabaseRemoteProcess'] && !options['DisableDatabase']
+      db_manager = Msf::DBManager.new(self)
+      db_manager.init_db(options)
+      options[:db_manager] = db_manager
+    end
+
+    Metasploit::Framework::DataService::DataProxy.new(options)
+  end
+
 end
 
 class FrameworkEventSubscriber

lib/msf/core/modules/metadata/store.rb

@@ -104,7 +104,7 @@ def configure_user_store
 
   def get_user_store
     store_dir = ::File.join(Msf::Config.config_directory, "store")
-    FileUtils.mkdir(store_dir) if !::File.exist?(store_dir)
+    FileUtils.makedirs(store_dir) if !::File.exist?(store_dir)
     return ::File.join(store_dir, UserMetaDataFile)
   end
 

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

@@ -393,7 +393,6 @@ def creds_search(*args)
     }
 
     tbl = Rex::Text::Table.new(tbl_opts)
-    opts = {}
     opts[:wspace] = framework.db.workspace
     query = framework.db.creds(opts)
 

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

@@ -90,13 +90,11 @@ def active?
           end
 
           def cmd_set_data_service(service_id)
-            data_proxy = Metasploit::Framework::DataService::DataProxy.instance
-            data_proxy.set_data_service(service_id)
+            framework.db.set_data_service(service_id)
           end
 
           def cmd_list_data_services()
-            data_service_manager = Metasploit::Framework::DataService::DataProxy.instance
-            data_service_manager.print_data_services
+            framework.db.print_data_services
           end
 
           def cmd_add_data_service(*args)
@@ -111,8 +109,7 @@ def cmd_add_data_service(*args)
 
             remote_service_endpoint = Metasploit::Framework::DataService::RemoteServiceEndpoint.new(host, port)
             remote_data_service = Metasploit::Framework::DataService::RemoteHTTPDataService.new(remote_service_endpoint)
-            data_service_manager = Metasploit::Framework::DataService::DataProxy.instance
-            data_service_manager.register_data_service(remote_data_service)
+            framework.db.register_data_service(remote_data_service)
           end
 
           def cmd_test_data_service_host(*args)
@@ -129,8 +126,7 @@ def cmd_test_data_service_host(*args)
             end
 
             puts 'Reporting test host to data service'
-            data_service = Metasploit::Framework::DataService::DataProxy.instance
-            data_service.report_host host
+            framework.db.report_host host
           end
 
           def cmd_test_data_service_loot(*args)
@@ -153,8 +149,7 @@ def cmd_test_data_service_loot(*args)
             end
 
             puts 'Reporting test loot to data service'
-            data_service = Metasploit::Framework::DataService::DataProxy.instance
-            data_service.report_loot loot
+            framework.db.report_loot loot
           end
 
           def cmd_perf_test_data_service_loot(*args)
@@ -198,11 +193,10 @@ def cmd_perf_test_data_service_loot(*args)
             end
 
             puts 'Reporting test loot to data service'
-            data_service = Metasploit::Framework::DataService::DataProxy.instance
             start_time = Time.now
             puts "#{start_time} - Staring loot perf test"
             loots.each do |loot|
-              data_service.report_loot loot
+              framework.db.report_loot loot
             end
             end_time = Time.now
             puts "#{end_time} - Ending loot perf test. Duration was #{end_time - start_time}"
@@ -256,9 +250,11 @@ def cmd_workspace(*args)
               end
               framework.db.workspace = workspace
             elsif deleting and names
-              framework.db.delete_workspaces(names)
+              status_msg, error_msg = framework.db.delete_workspaces(names)
+              print_msgs(status_msg, error_msg)
             elsif delete_all
-              framework.db.delete_all_workspaces()
+              status_msg, error_msg = framework.db.delete_all_workspaces()
+              print_msgs(status_msg, error_msg)
             elsif renaming
               if names.length != 2
                 print_error("Wrong number of arguments to rename")
@@ -1993,6 +1989,18 @@ def each_host_range_chunk(host_ranges, &block)
             end
           end
 
+          private
+
+          def print_msgs(status_msg, error_msg)
+            status_msg.each do |s|
+              print_status(s)
+            end
+
+            error_msg.each do |e|
+              print_error(e)
+            end
+          end
+
         end
 
       end end end end

lib/msf/ui/console/driver.rb

@@ -122,7 +122,6 @@ def initialize(prompt = DefaultPrompt, prompt_char = DefaultPromptChar, opts = {
       enstack_dispatcher(dispatcher)
     end
 
-    framework.db.init(framework, opts)
     if (framework.db.active)
         require 'msf/ui/console/command_dispatcher/db'
         enstack_dispatcher(CommandDispatcher::Db)

spec/lib/msf/core/auxiliary/cisco_spec.rb

@@ -77,7 +77,7 @@ def myworkspace
 
   context '#cisco_ios_config_eater' do
     before(:example) do
-      expect(aux_cisco).to receive(:myworkspace).and_return(workspace)
+      expect(aux_cisco).to receive(:myworkspace).at_least(:once).and_return(workspace)
     end
 
     it 'deals with udp ports' do