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."
|
raise Exception, "#{ui_message}: #{exception.message}. See log for more details."
|
||||||
end
|
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)
|
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.
|
# 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?
|
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.
|
# 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)
|
if opts[:workspace] && opts[:workspace].is_a?(::Mdm::Workspace)
|
||||||
opts[:workspace] = opts[:workspace].name
|
opts[:workspace] = opts[:workspace].name
|
||||||
|
|
|
@ -3,7 +3,7 @@ module WorkspaceDataProxy
|
||||||
def find_workspace(workspace_name)
|
def find_workspace(workspace_name)
|
||||||
begin
|
begin
|
||||||
data_service = self.get_data_service
|
data_service = self.get_data_service
|
||||||
opts = { :name => workspace_name }
|
opts = { name: workspace_name }
|
||||||
data_service.workspaces(opts).first
|
data_service.workspaces(opts).first
|
||||||
rescue Exception => e
|
rescue Exception => e
|
||||||
self.log_error(e, "Problem finding workspace")
|
self.log_error(e, "Problem finding workspace")
|
||||||
|
@ -38,7 +38,7 @@ module WorkspaceDataProxy
|
||||||
else
|
else
|
||||||
# This is mostly a failsafe to prevent bad things from happening. @current_workspace should always be set
|
# 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
|
# 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
|
@current_workspace = default_workspace
|
||||||
end
|
end
|
||||||
rescue Exception => e
|
rescue Exception => e
|
||||||
|
@ -46,8 +46,7 @@ module WorkspaceDataProxy
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
# TODO: Tracking of the current workspace should be moved out of the datastore.
|
# TODO: Tracking of the current workspace should be moved out of the datastore. See MS-3095.
|
||||||
# See MS-3095
|
|
||||||
def workspace=(workspace)
|
def workspace=(workspace)
|
||||||
begin
|
begin
|
||||||
@current_workspace = workspace
|
@current_workspace = workspace
|
||||||
|
@ -65,11 +64,9 @@ module WorkspaceDataProxy
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
def delete_workspaces(workspace_ids)
|
def delete_workspaces(opts)
|
||||||
begin
|
begin
|
||||||
data_service = self.get_data_service
|
data_service = self.get_data_service
|
||||||
opts = {}
|
|
||||||
opts[:ids] = workspace_ids
|
|
||||||
data_service.delete_workspaces(opts)
|
data_service.delete_workspaces(opts)
|
||||||
rescue Exception => e
|
rescue Exception => e
|
||||||
self.log_error(e, "Problem deleting workspaces")
|
self.log_error(e, "Problem deleting workspaces")
|
||||||
|
|
|
@ -122,12 +122,11 @@ class RemoteHTTPDataService
|
||||||
def make_request(request_type, path, data_hash = nil, query = nil)
|
def make_request(request_type, path, data_hash = nil, query = nil)
|
||||||
begin
|
begin
|
||||||
# simplify query by removing nil values
|
# simplify query by removing nil values
|
||||||
query_str = nil
|
|
||||||
query_str = (!query.nil? && !query.empty?) ? query.compact.to_query : nil
|
query_str = (!query.nil? && !query.empty?) ? query.compact.to_query : nil
|
||||||
uri = URI::HTTP::build({path: path, query: query_str})
|
uri = URI::HTTP::build({path: path, query: query_str})
|
||||||
dlog("HTTP #{request_type} request to #{uri.request_uri} with #{data_hash ? data_hash : "nil"}")
|
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
|
case request_type
|
||||||
when GET_REQUEST
|
when GET_REQUEST
|
||||||
request = Net::HTTP::Get.new(uri.request_uri)
|
request = Net::HTTP::Get.new(uri.request_uri)
|
||||||
|
@ -225,7 +224,7 @@ class RemoteHTTPDataService
|
||||||
raise 'Endpoint cannot be nil' if endpoint.nil?
|
raise 'Endpoint cannot be nil' if endpoint.nil?
|
||||||
end
|
end
|
||||||
|
|
||||||
def build_request(request, data_hash, add_workspace = true)
|
def build_request(request, data_hash)
|
||||||
request.content_type = 'application/json'
|
request.content_type = 'application/json'
|
||||||
if !data_hash.nil? && !data_hash.empty?
|
if !data_hash.nil? && !data_hash.empty?
|
||||||
data_hash.each do |k,v|
|
data_hash.each do |k,v|
|
||||||
|
|
|
@ -7,12 +7,12 @@ module RemoteWorkspaceDataService
|
||||||
WORKSPACE_MDM_CLASS = 'Mdm::Workspace'
|
WORKSPACE_MDM_CLASS = 'Mdm::Workspace'
|
||||||
|
|
||||||
def add_workspace(workspace_name)
|
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
|
json_to_mdm_object(response, WORKSPACE_MDM_CLASS, nil).first
|
||||||
end
|
end
|
||||||
|
|
||||||
def default_workspace
|
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
|
end
|
||||||
|
|
||||||
def workspace
|
def workspace
|
||||||
|
|
|
@ -154,7 +154,7 @@ class Db
|
||||||
names.each do |n|
|
names.each do |n|
|
||||||
ws_ids_to_delete << framework.db.find_workspace(n).id
|
ws_ids_to_delete << framework.db.find_workspace(n).id
|
||||||
end
|
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)
|
print_deleted_workspaces(deleted, starting_ws)
|
||||||
elsif delete_all
|
elsif delete_all
|
||||||
ws_ids_to_delete = []
|
ws_ids_to_delete = []
|
||||||
|
@ -162,7 +162,7 @@ class Db
|
||||||
framework.db.workspaces.each do |ws|
|
framework.db.workspaces.each do |ws|
|
||||||
ws_ids_to_delete << ws.id
|
ws_ids_to_delete << ws.id
|
||||||
end
|
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)
|
print_deleted_workspaces(deleted, starting_ws)
|
||||||
elsif renaming
|
elsif renaming
|
||||||
if names.length != 2
|
if names.length != 2
|
||||||
|
@ -221,10 +221,10 @@ class Db
|
||||||
ws.name,
|
ws.name,
|
||||||
framework.db.hosts(ws.name).count,
|
framework.db.hosts(ws.name).count,
|
||||||
framework.db.services(ws.name).count,
|
framework.db.services(ws.name).count,
|
||||||
framework.db.vulns({:workspace => ws.name}).count,
|
framework.db.vulns({workspace: ws.name}).count,
|
||||||
framework.db.creds({:workspace => ws.name}).count,
|
framework.db.creds({workspace: ws.name}).count,
|
||||||
framework.db.loots(ws.name).count,
|
framework.db.loots(ws.name).count,
|
||||||
framework.db.notes({:workspace => ws.name}).count
|
framework.db.notes({workspace: ws.name}).count
|
||||||
]
|
]
|
||||||
end
|
end
|
||||||
|
|
||||||
|
|
|
@ -21,6 +21,14 @@ module DBManager
|
||||||
condition_set.reduce { |conditions, condition| conditions.or(condition).expr }
|
condition_set.reduce { |conditions, condition| conditions.or(condition).expr }
|
||||||
end
|
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)
|
def self.process_opts_workspace(opts, framework, required = true)
|
||||||
wspace = delete_opts_workspace(opts)
|
wspace = delete_opts_workspace(opts)
|
||||||
if required && (wspace.nil? || ((wspace.kind_of? String) && wspace.empty?))
|
if required && (wspace.nil? || ((wspace.kind_of? String) && wspace.empty?))
|
||||||
|
@ -33,7 +41,12 @@ module DBManager
|
||||||
wspace
|
wspace
|
||||||
end
|
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)
|
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)
|
opts.delete(:workspace) || opts.delete(:wspace)
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
Loading…
Reference in New Issue