Finish moving current_workspace tracking to client

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

@@ -2,7 +2,8 @@ module EventDataProxy
 
   def report_event(opts)
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
+      opts[:workspace] = workspace.name if opts[:workspace].nil?
       data_service.report_event(opts)
     rescue  Exception => e
       self.log_error(e, "Problem reporting event")

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

@@ -29,10 +29,13 @@ module WorkspaceDataProxy
 
   def workspace
     begin
-      if @current_workspace_id
+      if @current_workspace
-        workspaces({ :id => @current_workspace_id }).first
+        @current_workspace
       else
-        default_workspace
+        # 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"
+        @current_workspace = default_workspace
       end
     rescue  Exception => e
       self.log_error(e, "Problem retrieving workspace")
@@ -43,7 +46,7 @@ module WorkspaceDataProxy
   # See MS-3095
   def workspace=(workspace)
     begin
-      @current_workspace_id = workspace.id
+      @current_workspace = workspace
     rescue  Exception => e
       self.log_error(e, "Problem setting workspace")
     end

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

@@ -24,10 +24,11 @@ class RemoteHTTPDataService
   #
   # @param [String] endpoint A valid http or https URL. Cannot be nil
   #
-  def initialize(endpoint, https_opts = {})
+  def initialize(endpoint, framework, https_opts = {})
     validate_endpoint(endpoint)
     @endpoint = URI.parse(endpoint)
     @https_opts = https_opts
+    @framework = framework
     build_client_pool(5)
   end
 
@@ -78,8 +79,8 @@ class RemoteHTTPDataService
   #
   # @return A wrapped response (ResponseWrapper), see below.
   #
-  def get_data(path, data_hash = nil, query = nil)
+  def get_data(path, data_hash = nil, query = nil, add_workspace = true)
-    make_request(GET_REQUEST, path, data_hash, query)
+    make_request(GET_REQUEST, path, data_hash, query, add_workspace)
   end
 
   #
@@ -118,10 +119,14 @@ class RemoteHTTPDataService
   #
   # @return A wrapped response (ResponseWrapper)
   #
-  def make_request(request_type, path, data_hash = nil, query = nil)
+  def make_request(request_type, path, data_hash = nil, query = nil, add_workspace = true)
     begin
       # simplify query by removing nil values
-      query_str = (!query.nil? && !query.empty?) ? append_workspace(query).compact.to_query : nil
+      add_workspace = true
+      query_str = nil
+      if add_workspace
+        query_str = (!query.nil? && !query.empty?) ? append_workspace(query).compact.to_query : nil
+      end
       uri = URI::HTTP::build({path: path, query: query_str})
       dlog("HTTP #{request_type} request to #{uri.request_uri} with #{data_hash ? data_hash : "nil"}")
 
@@ -227,12 +232,12 @@ class RemoteHTTPDataService
       data_hash[:workspace] = workspace.name
     end
 
-    data_hash[:workspace] = framework.db.workspace.name if workspace.nil?
+    data_hash[:workspace] = @framework.db.workspace.name if workspace.nil?
 
     data_hash
   end
 
-  def build_request(request, data_hash)
+  def build_request(request, data_hash, add_workspace = true)
     request.content_type = 'application/json'
     if !data_hash.nil? && !data_hash.empty?
       data_hash.each do |k,v|
@@ -244,7 +249,7 @@ class RemoteHTTPDataService
           data_hash.delete(k)
         end
       end
-      json_body = append_workspace(data_hash).to_json
+      json_body = append_workspace(data_hash).to_json if add_workspace
       request.body = json_body
     end
 

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

@@ -13,19 +13,25 @@ module RemoteWorkspaceDataService
   end
 
   def default_workspace
-    json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, nil, {:name => default}), WORKSPACE_MDM_CLASS, [])
+    json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, nil, {:name => default}, false), WORKSPACE_MDM_CLASS, [])
   end
 
   def workspace
-    find_workspace(current_workspace_name)
+    # The @current_workspace is tracked on the client side, so attempting to call it directly from the RemoteDataService
+    # will not return the correct results. Run it back through the proxy.
+    wlog "[DEPRECATION] Calling workspace from within the RemoteDataService is no longer supported. Please call from WorkspaceDataProxy instead."
+    framework.db.workspace
   end
 
   def workspace=(workspace)
-    @current_workspace_name = workspace.name
+    # The @current_workspace is tracked on the client side, so attempting to call it directly from the RemoteDataService
+    # will not return the correct results. Run it back through the proxy.
+    wlog "[DEPRECATION] Setting the current workspace from the RemoteDataService is no longer supported. Please call from WorkspaceDataProxy instead."
+    framework.db.workspace = workspace
   end
 
   def workspaces(opts)
-    json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, nil, opts), WORKSPACE_MDM_CLASS, [])
+    json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, nil, opts, false), WORKSPACE_MDM_CLASS, [])
   end
 
   def delete_workspaces(opts)

lib/msf/core/db_manager/workspace.rb

@@ -9,38 +9,38 @@ module Msf::DBManager::Workspace
   end
 
   def default_workspace
-    puts "default_workspace is being called directly from dbmanager"
+    # Workspace tracking is handled on the client side, so attempting to call it directly from the DbManager
-    caller.each { |line| puts "#{line}\n"}
+    # will not return the correct results. Run it back through the proxy.
-  # ::ActiveRecord::Base.connection_pool.with_connection {
+    wlog "[DEPRECATION] Setting the workspace from within DbManager is no longer supported. Please call from WorkspaceDataProxy instead."
-  #   ::Mdm::Workspace.default
+    raise NotImplementedError
-  # }
   end
 
   def find_workspace(name)
-    puts "find_workspace is being called directly from dbmanager"
+  ::ActiveRecord::Base.connection_pool.with_connection {
-    caller.each { |line| puts "#{line}\n"}
+    ::Mdm::Workspace.find_by_name(name)
-  # ::ActiveRecord::Base.connection_pool.with_connection {
+  }
-  #   ::Mdm::Workspace.find_by_name(name)
-  # }
   end
 
   def workspace
-    puts "workspace is being called directly from dbmanager"
+    # The @current_workspace is tracked on the client side, so attempting to call it directly from the DbManager
-    caller.each { |line| puts "#{line}\n"}
+    # will not return the correct results. Run it back through the proxy.
-  # ::ActiveRecord::Base.connection_pool.with_connection {
+    wlog "[DEPRECATION] Calling workspace from within DbManager is no longer supported. Please call from WorkspaceDataProxy instead."
-  #   ::Mdm::Workspace.find(@workspace_id)
+    raise NotImplementedError
-  # }
   end
 
   def workspace=(workspace)
-    #@workspace_id = workspace.id
+    # The @current_workspace is tracked on the client side, so attempting to call it directly from the DbManager
-    puts "workspace= is being called directly from dbmanager"
+    # will not return the correct results. Run it back through the proxy.
-    caller.each { |line| puts "#{line}\n"}
+    wlog "[DEPRECATION] Setting the workspace from within DbManager is no longer supported. Please call from WorkspaceDataProxy instead."
+    raise NotImplementedError
   end
 
   def workspaces(opts = {})
   ::ActiveRecord::Base.connection_pool.with_connection {
     search_term = opts.delete(:search_term)
+    # Passing these values to the search will cause exceptions, so remove them if they accidentally got passed in.
+    opts.delete(:workspace)
+    opts.delete(:wspace)
 
     ::ActiveRecord::Base.connection_pool.with_connection {
       if search_term && !search_term.empty?
@@ -58,7 +58,7 @@ module Msf::DBManager::Workspace
 
     ::ActiveRecord::Base.connection_pool.with_connection {
       deleted = []
-      default_deleted = false
+      default_deleted = false q
       opts[:ids].each do |ws_id|
         ws = Mdm::Workspace.find(ws_id)
         default_deleted = true if ws.default?

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

@@ -1963,7 +1963,7 @@ class Db
     end
 
     endpoint = "#{protocol}://#{host}:#{port}"
-    remote_data_service = Metasploit::Framework::DataService::RemoteHTTPDataService.new(endpoint, https_opts)
+    remote_data_service = Metasploit::Framework::DataService::RemoteHTTPDataService.new(endpoint, framework, https_opts)
     begin
       framework.db.register_data_service(remote_data_service)
       print_line "Registered data service: #{remote_data_service.name}"

lib/msf/ui/console/driver.rb

@@ -418,6 +418,7 @@ class Driver < Msf::Ui::Driver
         print_warning("\t#{path}: #{error}")
       end
     end
+    framework.db.workspace = framework.db.default_workspace
 
     framework.events.on_ui_start(Msf::Framework::Revision)
 

lib/msf/util/db_manager.rb

@@ -16,7 +16,9 @@ module DBManager
   end
 
   def self.process_opts_workspace(opts, framework)
-    wspace = opts.delete(:wspace) || opts.delete(:workspace) || framework.db.workspace
+    wspace = opts.delete(:wspace) || opts.delete(:workspace)
+    raise ArgumentError.new("opts must include a valid :workspace or :wspace value.") if wspace.nil? || ((wspace.kind_of? String) && wspace.empty?)
+
     if wspace.kind_of? String
       wspace = framework.db.find_workspace(wspace)
     end