Address PR suggestions and add comments
parent
68ad91763a
commit
f27490dc61
|
@ -123,12 +123,18 @@ class DataProxy
|
|||
raise Exception, "#{ui_message}: #{exception.message}. See log for more details."
|
||||
end
|
||||
|
||||
# Adds a valid workspace value to the opts hash before sending on to the data layer.
|
||||
#
|
||||
# @param [Hash] opts The opts hash that will be passed to the data layer.
|
||||
# @param [String] wspace A specific workspace name to add to the opts hash.
|
||||
# @return [Hash] The opts hash with a valid :workspace value added.
|
||||
def add_opts_workspace(opts, wspace = nil)
|
||||
opts[:workspace] = wspace if wspace
|
||||
|
||||
# Some methods use the key :wspace. Let's standardize on :workspace and clean it up here.
|
||||
opts[:workspace] = opts.delete(:wspace) unless opts[:wspace].nil?
|
||||
|
||||
# If the user passed in a specific workspace then use that in opts
|
||||
opts[:workspace] = wspace if wspace
|
||||
|
||||
# We only want to pass the workspace name, so grab it if it is currently an object.
|
||||
if opts[:workspace] && opts[:workspace].is_a?(::Mdm::Workspace)
|
||||
opts[:workspace] = opts[:workspace].name
|
||||
|
|
|
@ -3,7 +3,7 @@ module WorkspaceDataProxy
|
|||
def find_workspace(workspace_name)
|
||||
begin
|
||||
data_service = self.get_data_service
|
||||
opts = { :name => workspace_name }
|
||||
opts = { name: workspace_name }
|
||||
data_service.workspaces(opts).first
|
||||
rescue Exception => e
|
||||
self.log_error(e, "Problem finding workspace")
|
||||
|
@ -38,7 +38,7 @@ module WorkspaceDataProxy
|
|||
else
|
||||
# This is mostly a failsafe to prevent bad things from happening. @current_workspace should always be set
|
||||
# outside of here, but this will save us from crashes/infinite loops if that happens
|
||||
warn "@current_workspace was not set. Setting to default_workspace"
|
||||
warn "@current_workspace was not set. Setting to default_workspace: #{default_workspace.name}"
|
||||
@current_workspace = default_workspace
|
||||
end
|
||||
rescue Exception => e
|
||||
|
@ -46,8 +46,7 @@ module WorkspaceDataProxy
|
|||
end
|
||||
end
|
||||
|
||||
# TODO: Tracking of the current workspace should be moved out of the datastore.
|
||||
# See MS-3095
|
||||
# TODO: Tracking of the current workspace should be moved out of the datastore. See MS-3095.
|
||||
def workspace=(workspace)
|
||||
begin
|
||||
@current_workspace = workspace
|
||||
|
@ -65,11 +64,9 @@ module WorkspaceDataProxy
|
|||
end
|
||||
end
|
||||
|
||||
def delete_workspaces(workspace_ids)
|
||||
def delete_workspaces(opts)
|
||||
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")
|
||||
|
|
|
@ -122,12 +122,11 @@ class RemoteHTTPDataService
|
|||
def make_request(request_type, path, data_hash = nil, query = nil)
|
||||
begin
|
||||
# simplify query by removing nil values
|
||||
query_str = nil
|
||||
query_str = (!query.nil? && !query.empty?) ? query.compact.to_query : nil
|
||||
uri = URI::HTTP::build({path: path, query: query_str})
|
||||
dlog("HTTP #{request_type} request to #{uri.request_uri} with #{data_hash ? data_hash : "nil"}")
|
||||
|
||||
client = @client_pool.pop()
|
||||
client = @client_pool.pop
|
||||
case request_type
|
||||
when GET_REQUEST
|
||||
request = Net::HTTP::Get.new(uri.request_uri)
|
||||
|
@ -225,7 +224,7 @@ class RemoteHTTPDataService
|
|||
raise 'Endpoint cannot be nil' if endpoint.nil?
|
||||
end
|
||||
|
||||
def build_request(request, data_hash, add_workspace = true)
|
||||
def build_request(request, data_hash)
|
||||
request.content_type = 'application/json'
|
||||
if !data_hash.nil? && !data_hash.empty?
|
||||
data_hash.each do |k,v|
|
||||
|
|
|
@ -7,12 +7,12 @@ module RemoteWorkspaceDataService
|
|||
WORKSPACE_MDM_CLASS = 'Mdm::Workspace'
|
||||
|
||||
def add_workspace(workspace_name)
|
||||
response = self.post_data(WORKSPACE_API_PATH, {:workspace_name => workspace_name})
|
||||
response = self.post_data(WORKSPACE_API_PATH, { workspace_name: workspace_name })
|
||||
json_to_mdm_object(response, WORKSPACE_MDM_CLASS, nil).first
|
||||
end
|
||||
|
||||
def default_workspace
|
||||
json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, nil, {:name => Msf::DBManager::Workspace::DEFAULT_WORKSPACE_NAME}), WORKSPACE_MDM_CLASS, [])
|
||||
json_to_mdm_object(self.get_data(WORKSPACE_API_PATH, nil, { name: Msf::DBManager::Workspace::DEFAULT_WORKSPACE_NAME }), WORKSPACE_MDM_CLASS, [])
|
||||
end
|
||||
|
||||
def workspace
|
||||
|
|
|
@ -154,7 +154,7 @@ class Db
|
|||
names.each do |n|
|
||||
ws_ids_to_delete << framework.db.find_workspace(n).id
|
||||
end
|
||||
deleted = framework.db.delete_workspaces(ws_ids_to_delete)
|
||||
deleted = framework.db.delete_workspaces(ids: ws_ids_to_delete)
|
||||
print_deleted_workspaces(deleted, starting_ws)
|
||||
elsif delete_all
|
||||
ws_ids_to_delete = []
|
||||
|
@ -162,7 +162,7 @@ class Db
|
|||
framework.db.workspaces.each do |ws|
|
||||
ws_ids_to_delete << ws.id
|
||||
end
|
||||
deleted = framework.db.delete_workspaces(ws_ids_to_delete)
|
||||
deleted = framework.db.delete_workspaces(ids: ws_ids_to_delete)
|
||||
print_deleted_workspaces(deleted, starting_ws)
|
||||
elsif renaming
|
||||
if names.length != 2
|
||||
|
@ -221,10 +221,10 @@ class Db
|
|||
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.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
|
||||
framework.db.notes({workspace: ws.name}).count
|
||||
]
|
||||
end
|
||||
|
||||
|
|
|
@ -21,6 +21,14 @@ module DBManager
|
|||
condition_set.reduce { |conditions, condition| conditions.or(condition).expr }
|
||||
end
|
||||
|
||||
# Processes the workspace value in the opts hash from a request. This method throws an exception if
|
||||
# :workspace was not present but required was true, deletes the workspace from the hash, and
|
||||
# looks up the workspace object by name, which it returns.
|
||||
#
|
||||
# @param [Hash] opts The opts hash passed in from the data request. Should contain :workspace if required is true.
|
||||
# @param [Msf::Framework] framework A framework object containing a valid database connection.
|
||||
# @param [Bool] required true if the :workspace key is required for this data operation. false if it is only optional.
|
||||
# @return [Mdm::Workspace] The workspace object that was referenced by name in opts.
|
||||
def self.process_opts_workspace(opts, framework, required = true)
|
||||
wspace = delete_opts_workspace(opts)
|
||||
if required && (wspace.nil? || ((wspace.kind_of? String) && wspace.empty?))
|
||||
|
@ -33,7 +41,12 @@ module DBManager
|
|||
wspace
|
||||
end
|
||||
|
||||
# Removes the :workspace or :wspace key from the opts hash.
|
||||
#
|
||||
# @param [Hash] opts The opts hash passed in from the data request.
|
||||
# @return [String] The name of the workspace that was contained in the key.
|
||||
def self.delete_opts_workspace(opts)
|
||||
wlog("Both :workspace and :wspace were found in opts. Using :workspace.") if opts[:workspace] && opts[:wspace]
|
||||
opts.delete(:workspace) || opts.delete(:wspace)
|
||||
end
|
||||
end
|
||||
|
|
Loading…
Reference in New Issue