Land #5902, workspace -D all the workspaces

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

@@ -89,6 +89,7 @@ class Db
     print_line "    workspace [name]           Switch workspace"
     print_line "    workspace -a [name] ...    Add workspace(s)"
     print_line "    workspace -d [name] ...    Delete workspace(s)"
+    print_line "    workspace -D               Delete all workspaces"
     print_line "    workspace -r <old> <new>   Rename workspace"
     print_line "    workspace -h               Show this help information"
     print_line
@@ -106,6 +107,8 @@ class Db
         adding = true
       when '-d','--del'
         deleting = true
+      when '-D','--delete-all'
+        delete_all = true
       when '-r','--rename'
         renaming = true
       else
@@ -123,28 +126,9 @@ class Db
       end
       framework.db.workspace = workspace
     elsif deleting and names
-      switched = false
+      delete_workspaces(names)
-      # Delete workspaces
+    elsif delete_all
-      names.each do |name|
+      delete_workspaces(framework.db.workspaces.map(&:name))
-        workspace = framework.db.find_workspace(name)
-        if workspace.nil?
-          print_error("Workspace not found: #{name}")
-        elsif workspace.default?
-          workspace.destroy
-          workspace = framework.db.add_workspace(name)
-          print_status("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
-            framework.db.workspace = framework.db.default_workspace
-            switched = true
-          end
-          # now destroy the named workspace
-          workspace.destroy
-          print_status("Deleted workspace: #{name}")
-        end
-      end
-      print_status("Switched workspace: #{framework.db.workspace.name}") if switched
     elsif renaming
       if names.length != 2
         print_error("Wrong number of arguments to rename")
@@ -202,6 +186,31 @@ class Db
   }
   end
 
+  def delete_workspaces(names)
+    switched = false
+    # Delete workspaces
+    names.each do |name|
+      workspace = framework.db.find_workspace(name)
+      if workspace.nil?
+        print_error("Workspace not found: #{name}")
+      elsif workspace.default?
+        workspace.destroy
+        workspace = framework.db.add_workspace(name)
+        print_status("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
+          framework.db.workspace = framework.db.default_workspace
+          switched = true
+        end
+        # now destroy the named workspace
+        workspace.destroy
+        print_status("Deleted workspace: #{name}")
+      end
+    end
+    print_status("Switched workspace: #{framework.db.workspace.name}") if switched
+  end
+
   def cmd_workspace_tabs(str, words)
     return [] unless active?
     framework.db.workspaces.map { |s| s.name } if (words & ['-a','--add']).empty?

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

@@ -535,6 +535,65 @@ describe Msf::Ui::Console::CommandDispatcher::Db do
   end
 
   describe "#cmd_workspace" do
+    before(:each) do
+      db.cmd_workspace "-D"
+      @output = []
+    end
+    describe "<no arguments>" do
+      it "should list default workspace" do
+        db.cmd_workspace
+        @output.should =~ [
+          "* default"
+        ]
+      end
+
+      it "should list all workspaces" do
+        db.cmd_workspace("-a", "foo")
+        @output = []
+        db.cmd_workspace
+        @output.should =~ [
+          "  default",
+          "* foo"
+        ]
+      end
+    end
+
+    describe "-a" do
+      it "should add workspaces" do
+        db.cmd_workspace("-a", "foo", "bar", "baf")
+        @output.should =~ [
+          "Added workspace: foo",
+          "Added workspace: bar",
+          "Added workspace: baf"
+        ]
+      end
+    end
+
+    describe "-d" do
+      it "should delete a workspace" do
+        db.cmd_workspace("-a", "foo")
+        @output = []
+        db.cmd_workspace("-d", "foo")
+        @output.should =~ [
+          "Deleted workspace: foo",
+          "Switched workspace: default"
+        ]
+      end
+    end
+
+    describe "-D" do
+      it "should delete all workspaces" do
+        db.cmd_workspace("-a", "foo")
+        @output = []
+        db.cmd_workspace("-D")
+        @output.should =~ [
+          "Deleted and recreated the default workspace",
+          "Deleted workspace: foo",
+          "Switched workspace: default"
+        ]
+      end
+    end
+
     describe "-h" do
       it "should show a help message" do
         db.cmd_workspace "-h"
@@ -544,6 +603,7 @@ describe Msf::Ui::Console::CommandDispatcher::Db do
           "    workspace [name]           Switch workspace",
           "    workspace -a [name] ...    Add workspace(s)",
           "    workspace -d [name] ...    Delete workspace(s)",
+          "    workspace -D               Delete all workspaces",
           "    workspace -r <old> <new>   Rename workspace",
           "    workspace -h               Show this help information"
         ]

spec/support/shared/contexts/msf/ui_driver.rb

@@ -9,6 +9,10 @@ shared_context 'Msf::UIDriver' do
         @output ||= []
         @output.concat string.split("\n")
       end
+      driver.stub(:print_status).with(kind_of(String)) do |string|
+        @output ||= []
+        @output.concat string.split("\n")
+      end
       driver.stub(:print_error).with(kind_of(String)) do |string|
         @error ||= []
         @error.concat string.split("\n")