diff --git a/lib/msf/ui/console/command_dispatcher/db.rb b/lib/msf/ui/console/command_dispatcher/db.rb index 13287d2d36..1db20e3aee 100644 --- a/lib/msf/ui/console/command_dispatcher/db.rb +++ b/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 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.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 + delete_workspaces(names) + elsif delete_all + delete_workspaces(framework.db.workspaces.map(&:name)) 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? diff --git a/spec/lib/msf/ui/console/command_dispatcher/db_spec.rb b/spec/lib/msf/ui/console/command_dispatcher/db_spec.rb index ef35238716..dc1f75f6b6 100644 --- a/spec/lib/msf/ui/console/command_dispatcher/db_spec.rb +++ b/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 "" 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 Rename workspace", " workspace -h Show this help information" ] diff --git a/spec/support/shared/contexts/msf/ui_driver.rb b/spec/support/shared/contexts/msf/ui_driver.rb index 985914246a..2d64576574 100644 --- a/spec/support/shared/contexts/msf/ui_driver.rb +++ b/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")