Refactor workspace delete to be consistent with other commands

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

@@ -2,7 +2,7 @@ module WorkspaceDataProxy
 
   def find_workspace(workspace_name)
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
       data_service.find_workspace(workspace_name)
     rescue  Exception => e
       self.log_error(e, "Problem finding workspace")
@@ -11,7 +11,7 @@ module WorkspaceDataProxy
 
   def add_workspace(workspace_name)
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
       data_service.add_workspace(workspace_name)
     rescue  Exception => e
       self.log_error(e, "Problem adding workspace")
@@ -20,7 +20,7 @@ module WorkspaceDataProxy
 
   def default_workspace
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
       data_service.default_workspace
     rescue  Exception => e
       self.log_error(e, "Problem finding default workspace")
@@ -29,7 +29,7 @@ module WorkspaceDataProxy
 
   def workspace
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
       data_service.workspace
     rescue  Exception => e
       self.log_error(e, "Problem retrieving workspace")
@@ -38,7 +38,7 @@ module WorkspaceDataProxy
 
   def workspace=(workspace)
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
       data_service.workspace = workspace
     rescue  Exception => e
       self.log_error(e, "Problem setting workspace")
@@ -47,10 +47,21 @@ module WorkspaceDataProxy
 
   def workspaces
     begin
-      data_service = self.get_data_service()
+      data_service = self.get_data_service
       data_service.workspaces
     rescue  Exception => e
       self.log_error(e, "Problem retrieving workspaces")
     end
   end
+
+  def delete_workspaces(workspace_ids)
+    begin
+      data_service = self.get_data_service
+      opts = {}
+      opts[:ids] = workspace_ids
+      data_service.delete_workspaces(opts)
+    rescue Exception => e
+      self.log_error(e, "Problem deleting workspaces")
+    end
+  end
 end

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

@@ -3,8 +3,6 @@ require 'metasploit/framework/data_service/remote/http/response_data_helper'
 module RemoteWorkspaceDataService
   include ResponseDataHelper
 
-  # TODO: should counts be a flag in query data for the workspaces resource?
-  WORKSPACE_COUNTS_API_PATH = '/api/v1/workspaces/counts'
   WORKSPACE_API_PATH = '/api/v1/workspaces'
   WORKSPACE_MDM_CLASS = 'Mdm::Workspace'
   DEFAULT_WORKSPACE_NAME = 'default'
@@ -38,6 +36,10 @@ module RemoteWorkspaceDataService
     json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, {:all => true}), WORKSPACE_MDM_CLASS, [])
   end
 
+  def delete_workspaces(opts)
+    json_to_mdm_object(self.delete_data(WORKSPACE_API_PATH, opts), WORKSPACE_MDM_CLASS, [])
+  end
+
   #########
   protected
   #########
@@ -50,4 +52,4 @@ module RemoteWorkspaceDataService
     @current_workspace_name ||= DEFAULT_WORKSPACE_NAME
   end
 
-end
+end

lib/msf/core/db_manager/workspace.rb

@@ -34,33 +34,29 @@ module Msf::DBManager::Workspace
   }
   end
 
-  def delete_workspaces(names)
-    status_msg = []
-    error_msg = []
+  def delete_workspaces(opts)
+    raise ArgumentError.new("The following options are required: :ids") if opts[:ids].nil?
 
-    switched = false
-    # Delete workspaces
-    names.each do |name|
-      workspace = framework.db.find_workspace(name)
-      if workspace.nil?
-        error << "Workspace not found: #{name}"
-      elsif workspace.default?
-        workspace.destroy
-        workspace = framework.db.add_workspace(name)
-        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
+    ::ActiveRecord::Base.connection_pool.with_connection {
+      deleted = []
+      default_deleted = false
+      opts[:ids].each do |ws_id|
+        ws = Mdm::Workspace.find(ws_id)
+        default_deleted = true if ws.default?
+        if framework.db.workspace.name == ws.name
           framework.db.workspace = framework.db.default_workspace
-          switched = true
         end
-        # now destroy the named workspace
-        workspace.destroy
-        status_msg << "Deleted workspace: #{name}"
+        begin
+          deleted << ws.destroy
+          framework.db.workspace = framework.db.add_workspace('default') if default_deleted
+        rescue
+          elog("Forcibly deleting #{workspace}")
+          deleted << ws.delete
+        end
       end
-    end
-    (status_msg << "Switched workspace: #{framework.db.workspace.name}") if switched
-    return status_msg, error_msg
+
+      return deleted
+    }
   end
 
   #

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

@@ -149,11 +149,21 @@ class Db
       end
       framework.db.workspace = workspace
     elsif deleting and names
-      status_msg, error_msg = framework.db.delete_workspaces(names)
-      print_msgs(status_msg, error_msg)
+      ws_ids_to_delete = []
+      starting_ws = framework.db.workspace
+      names.each do |n|
+        ws_ids_to_delete << framework.db.find_workspace(n).id
+      end
+      deleted = framework.db.delete_workspaces(ws_ids_to_delete)
+      print_deleted_workspaces(deleted, starting_ws)
     elsif delete_all
-      status_msg, error_msg = framework.db.delete_all_workspaces()
-      print_msgs(status_msg, error_msg)
+      ws_ids_to_delete = []
+      starting_ws = framework.db.workspace
+      framework.db.workspaces.each do |ws|
+        ws_ids_to_delete << ws.id
+      end
+      deleted = framework.db.delete_workspaces(ws_ids_to_delete)
+      print_deleted_workspaces(deleted, starting_ws)
     elsif renaming
       if names.length != 2
         print_error("Wrong number of arguments to rename")
@@ -218,6 +228,18 @@ class Db
     end
   end
 
+  def print_deleted_workspaces(deleted_workspaces, starting_ws)
+    deleted_workspaces.each do |ws|
+      if ws.name == 'default'
+        print_status 'Deleted and recreated the default workspace'
+      elsif ws == starting_ws
+        print_status "Switched workspace: #{framework.db.workspace.name}"
+      else
+        print_status "Deleted workspace: #{ws.name}"
+      end
+    end
+  end
+
   def cmd_workspace_tabs(str, words)
     return [] unless active?
     framework.db.workspaces.map { |s| s.name } if (words & ['-a','--add']).empty?