Address PR suggestions and add comments

GSoC/Meterpreter_Web_Console
James Barnett 2018-04-16 16:45:23 -05:00
parent 68ad91763a
commit f27490dc61
No known key found for this signature in database
GPG Key ID: 647983861A4EC5EA
6 changed files with 34 additions and 19 deletions

View File

@ -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

View File

@ -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")

View File

@ -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|

View File

@ -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

View File

@ -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

View File

@ -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