Skip to content

Commit

Permalink
Consolidate logic for actions taken on a workspace
Browse files Browse the repository at this point in the history
* When renaming a workspace allow validation to handle all cases the same
* Do not display backtrace on console for rename failure, this is still logged to framework.log
  • Loading branch information
jmartin-tech committed Oct 15, 2019
1 parent b1ae404 commit 7505e1b
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 8 deletions.
11 changes: 8 additions & 3 deletions lib/msf/core/db_manager/workspace.rb
Original file line number Diff line number Diff line change
Expand Up @@ -97,11 +97,16 @@ def update_workspace(opts)
Msf::Util::DBManager.delete_opts_workspace(opts)

::ActiveRecord::Base.connection_pool.with_connection {
ws_to_update = workspaces({ id: opts.delete(:id) }).first
ws_id = opts.delete(:id)
ws_to_update = workspaces({ id: ws_id }).first
default_renamed = true if ws_to_update.name == DEFAULT_WORKSPACE_NAME
updated_ws = ws_to_update.update!(opts)
begin
ws_to_update.update!(opts) # will raise exception if it fails
rescue ActiveRecord::RecordInvalid => e
raise ArgumentError.new(e.message)
end
add_workspace({ name: DEFAULT_WORKSPACE_NAME }) if default_renamed
updated_ws
workspaces({ id: ws_id }).first
}
end
end
5 changes: 0 additions & 5 deletions lib/msf/ui/console/command_dispatcher/db.rb
Original file line number Diff line number Diff line change
Expand Up @@ -179,10 +179,6 @@ def cmd_workspace(*args)
name: names.last
}
begin
if names.last == Msf::DBManager::Workspace::DEFAULT_WORKSPACE_NAME
print_error("Unable to rename a workspace to '#{Msf::DBManager::Workspace::DEFAULT_WORKSPACE_NAME}'")
return
end
updated_ws = framework.db.update_workspace(opts)
if updated_ws
framework.db.workspace = updated_ws if names.first == framework.db.workspace.name
Expand All @@ -197,7 +193,6 @@ def cmd_workspace(*args)
end
rescue => e
print_error "Failed to rename workspace: #{e.message}"
e.backtrace.each { |line| print_error "#{line}"}
end

elsif names
Expand Down

0 comments on commit 7505e1b

Please sign in to comment.