From e5513409db11f22a68f8dc8f2671fc55fd2cc515 Mon Sep 17 00:00:00 2001
From: James Barnett <James_Barnett@rapid7.com>
Date: Tue, 24 Apr 2018 13:53:55 -0500
Subject: [PATCH] Include :workspace in db_import opts

---
 .../proxy/data_proxy_auto_loader.rb           |  2 ++
 .../proxy/db_import_data_proxy.rb             | 21 +++++++++++++++++++
 .../data_service/stubs/db_export_service.rb   |  5 +++++
 .../data_service/stubs/db_import_service.rb   |  9 ++++++++
 lib/msf/core/db_manager/import.rb             | 10 ++++-----
 5 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 lib/metasploit/framework/data_service/proxy/db_import_data_proxy.rb
 create mode 100644 lib/metasploit/framework/data_service/stubs/db_export_service.rb
 create mode 100644 lib/metasploit/framework/data_service/stubs/db_import_service.rb

diff --git a/lib/metasploit/framework/data_service/proxy/data_proxy_auto_loader.rb b/lib/metasploit/framework/data_service/proxy/data_proxy_auto_loader.rb
index ee6ef1949a..ce0eda6211 100644
--- a/lib/metasploit/framework/data_service/proxy/data_proxy_auto_loader.rb
+++ b/lib/metasploit/framework/data_service/proxy/data_proxy_auto_loader.rb
@@ -17,6 +17,7 @@ module DataProxyAutoLoader
   autoload :CredentialDataProxy, 'metasploit/framework/data_service/proxy/credential_data_proxy'
   autoload :NmapDataProxy, 'metasploit/framework/data_service/proxy/nmap_data_proxy'
   autoload :DbExportDataProxy, 'metasploit/framework/data_service/proxy/db_export_data_proxy'
+  autoload :DbImportDataProxy, 'metasploit/framework/data_service/proxy/db_import_data_proxy'
   autoload :VulnAttemptDataProxy, 'metasploit/framework/data_service/proxy/vuln_attempt_data_proxy'
 
   include ServiceDataProxy
@@ -33,5 +34,6 @@ module DataProxyAutoLoader
   include CredentialDataProxy
   include NmapDataProxy
   include DbExportDataProxy
+  include DbImportDataProxy
   include VulnAttemptDataProxy
 end
\ No newline at end of file
diff --git a/lib/metasploit/framework/data_service/proxy/db_import_data_proxy.rb b/lib/metasploit/framework/data_service/proxy/db_import_data_proxy.rb
new file mode 100644
index 0000000000..5e4ae64b44
--- /dev/null
+++ b/lib/metasploit/framework/data_service/proxy/db_import_data_proxy.rb
@@ -0,0 +1,21 @@
+module DbImportDataProxy
+  def import(opts, &block)
+    begin
+      data_service = self.get_data_service
+      add_opts_workspace(opts)
+      data_service.import(opts, &block)
+    rescue Exception => e
+      self.log_error(e, "Problem generating DB Export")
+    end
+  end
+
+  def import_file(opts, &block)
+    begin
+      data_service = self.get_data_service
+      add_opts_workspace(opts)
+      data_service.import_file(opts, &block)
+    rescue Exception => e
+      self.log_error(e, "Problem generating DB Export")
+    end
+  end
+end
diff --git a/lib/metasploit/framework/data_service/stubs/db_export_service.rb b/lib/metasploit/framework/data_service/stubs/db_export_service.rb
new file mode 100644
index 0000000000..90164573d3
--- /dev/null
+++ b/lib/metasploit/framework/data_service/stubs/db_export_service.rb
@@ -0,0 +1,5 @@
+module DbExportDataService
+  def run_db_export(path, format)
+    raise 'DbExportDataService#run_db_export is not implemented'
+  end
+end
\ No newline at end of file
diff --git a/lib/metasploit/framework/data_service/stubs/db_import_service.rb b/lib/metasploit/framework/data_service/stubs/db_import_service.rb
new file mode 100644
index 0000000000..952fde000c
--- /dev/null
+++ b/lib/metasploit/framework/data_service/stubs/db_import_service.rb
@@ -0,0 +1,9 @@
+module DbImportDataService
+  def import(opts, &block)
+    raise 'DbImportDataService#import is not implemented'
+  end
+
+  def import_file(opts, &block)
+    raise 'DbImportDataService#import_file is not implemented'
+  end
+end
\ No newline at end of file
diff --git a/lib/msf/core/db_manager/import.rb b/lib/msf/core/db_manager/import.rb
index be99b567ca..61702f442d 100644
--- a/lib/msf/core/db_manager/import.rb
+++ b/lib/msf/core/db_manager/import.rb
@@ -85,14 +85,14 @@ module Msf::DBManager::Import
   # import_file_detect will raise an error if the filetype
   # is unknown.
   def import(args={}, &block)
-    wspace = args[:wspace] || args['wspace'] || workspace
+    wspace = Msf::Util::DBManager.process_opts_workspace(args, framework)
     preserve_hosts = args[:task].options["DS_PRESERVE_HOSTS"] if args[:task].present? && args[:task].options.present?
     wspace.update_attribute(:import_fingerprint, true)
     existing_host_ids = wspace.hosts.map(&:id)
     data = args[:data] || args['data']
     ftype = import_filetype_detect(data)
     yield(:filetype, @import_filedata[:type]) if block
-    self.send "import_#{ftype}".to_sym, args, &block
+    self.send "import_#{ftype}".to_sym, args.merge(workspace: wspace.name), &block
     if preserve_hosts
       new_host_ids = Mdm::Host.where(workspace: wspace).map(&:id)
       (new_host_ids - existing_host_ids).each do |id|
@@ -111,7 +111,7 @@ module Msf::DBManager::Import
   #
   def import_file(args={}, &block)
     filename = args[:filename] || args['filename']
-    wspace = args[:wspace] || args['wspace'] || workspace
+    wspace = Msf::Util::DBManager.process_opts_workspace(args, framework)
     @import_filedata            = {}
     @import_filedata[:filename] = filename
 
@@ -148,9 +148,9 @@ module Msf::DBManager::Import
     REXML::Security.entity_expansion_text_limit = 51200
 
     if block
-      import(args.merge(:data => data)) { |type,data| yield type,data }
+      import(args.merge(data: data, workspace: wspace.name)) { |type,data| yield type,data }
     else
-      import(args.merge(:data => data))
+      import(args.merge(data: data, workspace: wspace.name))
     end
   end