From d23e33a7bc424511e9736d9ce941d7c81e7f9234 Mon Sep 17 00:00:00 2001 From: James Barnett Date: Wed, 28 Mar 2018 13:31:49 -0500 Subject: [PATCH] Make workspace -v work with updates --- .../framework/data_service/proxy/core.rb | 5 ++++ .../data_service/proxy/event_data_proxy.rb | 2 +- .../data_service/proxy/exploit_data_proxy.rb | 9 ++++--- .../data_service/proxy/host_data_proxy.rb | 14 +++++++---- .../data_service/proxy/loot_data_proxy.rb | 4 ++- .../data_service/proxy/service_data_proxy.rb | 11 +++++--- .../data_service/remote/http/core.rb | 4 +-- lib/msf/ui/console/command_dispatcher/db.rb | 25 ++++++++----------- lib/msf/util/db_manager.rb | 2 +- 9 files changed, 45 insertions(+), 31 deletions(-) diff --git a/lib/metasploit/framework/data_service/proxy/core.rb b/lib/metasploit/framework/data_service/proxy/core.rb index e03db7cad1..b2ea44b151 100644 --- a/lib/metasploit/framework/data_service/proxy/core.rb +++ b/lib/metasploit/framework/data_service/proxy/core.rb @@ -123,6 +123,11 @@ class DataProxy raise Exception, "#{ui_message}: #{exception.message}. See log for more details." end + def add_opts_workspace(opts) + opts[:workspace] = workspace if opts[:workspace].nil? && opts[:wspace].nil? + opts + end + ####### private ####### diff --git a/lib/metasploit/framework/data_service/proxy/event_data_proxy.rb b/lib/metasploit/framework/data_service/proxy/event_data_proxy.rb index a465fd84f7..d10c188772 100644 --- a/lib/metasploit/framework/data_service/proxy/event_data_proxy.rb +++ b/lib/metasploit/framework/data_service/proxy/event_data_proxy.rb @@ -3,7 +3,7 @@ module EventDataProxy def report_event(opts) begin data_service = self.get_data_service - opts[:workspace] = workspace.name if opts[:workspace].nil? + add_opts_workspace(opts) data_service.report_event(opts) rescue Exception => e self.log_error(e, "Problem reporting event") diff --git a/lib/metasploit/framework/data_service/proxy/exploit_data_proxy.rb b/lib/metasploit/framework/data_service/proxy/exploit_data_proxy.rb index f3f7817c2d..3d5fabed9a 100644 --- a/lib/metasploit/framework/data_service/proxy/exploit_data_proxy.rb +++ b/lib/metasploit/framework/data_service/proxy/exploit_data_proxy.rb @@ -2,7 +2,8 @@ module ExploitDataProxy def report_exploit_attempt(host, opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.report_exploit_attempt(host, opts) rescue Exception => e self.log_error(e, "Problem reporting exploit attempt") @@ -11,7 +12,8 @@ module ExploitDataProxy def report_exploit_failure(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.report_exploit_failure(opts) rescue Exception => e self.log_error(e, "Problem reporting exploit failure") @@ -20,7 +22,8 @@ module ExploitDataProxy def report_exploit_success(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.report_exploit_success(opts) rescue Exception => e self.log_error(e, "Problem reporting exploit success") diff --git a/lib/metasploit/framework/data_service/proxy/host_data_proxy.rb b/lib/metasploit/framework/data_service/proxy/host_data_proxy.rb index 196838d8b5..cf7b2d3f53 100644 --- a/lib/metasploit/framework/data_service/proxy/host_data_proxy.rb +++ b/lib/metasploit/framework/data_service/proxy/host_data_proxy.rb @@ -2,7 +2,7 @@ module HostDataProxy def hosts(wspace = workspace, non_dead = false, addresses = nil, search_term = nil) begin - data_service = self.get_data_service() + data_service = self.get_data_service opts = {} opts[:wspace] = wspace opts[:non_dead] = non_dead @@ -24,7 +24,8 @@ module HostDataProxy return unless valid(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.report_host(opts) rescue Exception => e self.log_error(e, "Problem reporting host") @@ -33,7 +34,8 @@ module HostDataProxy def report_hosts(hosts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.report_hosts(hosts) rescue Exception => e self.log_error(e, "Problem reporting hosts") @@ -42,7 +44,8 @@ module HostDataProxy def update_host(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.update_host(opts) rescue Exception => e self.log_error(e, "Problem updating host") @@ -51,7 +54,8 @@ module HostDataProxy def delete_host(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.delete_host(opts) rescue Exception => e self.log_error(e, "Problem deleting host") diff --git a/lib/metasploit/framework/data_service/proxy/loot_data_proxy.rb b/lib/metasploit/framework/data_service/proxy/loot_data_proxy.rb index 9c4d35b68c..305830e839 100644 --- a/lib/metasploit/framework/data_service/proxy/loot_data_proxy.rb +++ b/lib/metasploit/framework/data_service/proxy/loot_data_proxy.rb @@ -2,10 +2,11 @@ module LootDataProxy def report_loot(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service if !data_service.is_a?(Msf::DBManager) opts[:data] = Base64.urlsafe_encode64(opts[:data]) if opts[:data] end + add_opts_workspace(opts) data_service.report_loot(opts) rescue Exception => e self.log_error(e, "Problem reporting loot") @@ -33,6 +34,7 @@ module LootDataProxy def update_loot(opts) begin data_service = self.get_data_service + add_opts_workspace(opts) data_service.update_loot(opts) rescue Exception => e self.log_error(e, "Problem updating loot") diff --git a/lib/metasploit/framework/data_service/proxy/service_data_proxy.rb b/lib/metasploit/framework/data_service/proxy/service_data_proxy.rb index cfe6f14d83..dfb630eefb 100644 --- a/lib/metasploit/framework/data_service/proxy/service_data_proxy.rb +++ b/lib/metasploit/framework/data_service/proxy/service_data_proxy.rb @@ -2,7 +2,7 @@ module ServiceDataProxy def services(wspace = workspace, opts = {}) begin - data_service = self.get_data_service() + data_service = self.get_data_service opts[:workspace] = wspace data_service.services(opts) rescue Exception => e @@ -16,7 +16,8 @@ module ServiceDataProxy def report_service(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.report_service(opts) rescue Exception => e self.log_error(e, 'Problem reporting service') @@ -25,7 +26,8 @@ module ServiceDataProxy def update_service(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.update_service(opts) rescue Exception => e self.log_error(e, 'Problem updating service') @@ -34,7 +36,8 @@ module ServiceDataProxy def delete_service(opts) begin - data_service = self.get_data_service() + data_service = self.get_data_service + add_opts_workspace(opts) data_service.delete_service(opts) rescue Exception => e self.log_error(e, 'Problem deleting service') diff --git a/lib/metasploit/framework/data_service/remote/http/core.rb b/lib/metasploit/framework/data_service/remote/http/core.rb index e8c3ca4677..34bba4f623 100644 --- a/lib/metasploit/framework/data_service/remote/http/core.rb +++ b/lib/metasploit/framework/data_service/remote/http/core.rb @@ -225,8 +225,8 @@ class RemoteHTTPDataService end def append_workspace(data_hash) - workspace = data_hash[:workspace] - workspace = data_hash.delete(:wspace) unless workspace + # Some methods use the key :wspace. Let's standardize on :workspace and clean it up here. + workspace = data_hash.keys.include?(:wspace) ? data_hash.delete(:wspace) : data_hash[:workspace] if workspace && (workspace.is_a?(OpenStruct) || workspace.is_a?(::Mdm::Workspace)) data_hash[:workspace] = workspace.name diff --git a/lib/msf/ui/console/command_dispatcher/db.rb b/lib/msf/ui/console/command_dispatcher/db.rb index 8ffc336414..38a962c2b7 100644 --- a/lib/msf/ui/console/command_dispatcher/db.rb +++ b/lib/msf/ui/console/command_dispatcher/db.rb @@ -192,12 +192,12 @@ class Db return end else - workspace = framework.db.workspace + current_workspace = framework.db.workspace unless verbose current = nil framework.db.workspaces.sort_by {|s| s.name}.each do |s| - if s.name == workspace.name + if s.name == current_workspace.name current = s.name else print_line(" #{s.name}") @@ -206,8 +206,6 @@ class Db print_line("%red* #{current}%clr") unless current.nil? return end - workspace = framework.db.workspace - col_names = %w{current name hosts services vulns creds loots notes} tbl = Rex::Text::Table.new( @@ -217,17 +215,16 @@ class Db 'SearchTerm' => search_term ) - # List workspaces - framework.db.workspace_associations_counts.each do |ws| + framework.db.workspaces.each do |ws| tbl << [ - ws[:name] == workspace.name ? '*' : '', - ws[:name], - ws[:hosts_count], - ws[:services_count], - ws[:vulns_count], - ws[:creds_count], - ws[:loots_count], - ws[:notes_count] + current_workspace.name == ws.name ? '*' : '', + ws.name, + framework.db.hosts(ws.name).count, + framework.db.services(ws.name).count, + framework.db.vulns({:workspace => ws.name}).count, + framework.db.creds({:workspace => ws.name}).count, + framework.db.loots(ws.name).count, + framework.db.notes({:workspace => ws.name}).count ] end diff --git a/lib/msf/util/db_manager.rb b/lib/msf/util/db_manager.rb index 692b5f426b..ec912df462 100644 --- a/lib/msf/util/db_manager.rb +++ b/lib/msf/util/db_manager.rb @@ -17,7 +17,7 @@ module DBManager def self.process_opts_workspace(opts, framework) 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?) + raise ArgumentError.new("opts must include a valid :workspace.") if wspace.nil? || ((wspace.kind_of? String) && wspace.empty?) if wspace.kind_of? String wspace = framework.db.find_workspace(wspace)