Add with_connection wrappers to the database rpc calls
Certainly not all of these methods require a connection, but it is better to check one out when we don't need it than to risk grabbing an implicit connection that will never be handed back to the pool.unstable
parent
c68a775106
commit
d79f8b0492
|
@ -23,6 +23,7 @@ private
|
|||
end
|
||||
|
||||
def opts_to_hosts(opts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
wspace = find_workspace(opts[:workspace])
|
||||
hosts = []
|
||||
if opts[:host] or opts[:address]
|
||||
|
@ -39,9 +40,11 @@ private
|
|||
hosts |= hent if hent.class == Array
|
||||
end
|
||||
return hosts
|
||||
}
|
||||
end
|
||||
|
||||
def opts_to_services(hosts,opts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
wspace = find_workspace(opts[:workspace])
|
||||
services = []
|
||||
if opts[:host] or opts[:address] or opts[:addresses]
|
||||
|
@ -68,6 +71,7 @@ private
|
|||
services << sret if sret.class == ::Mdm::Service
|
||||
end
|
||||
return services
|
||||
}
|
||||
end
|
||||
|
||||
def db_check
|
||||
|
@ -84,6 +88,7 @@ private
|
|||
public
|
||||
|
||||
def rpc_hosts(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
|
||||
conditions = {}
|
||||
|
@ -113,9 +118,11 @@ public
|
|||
ret[:hosts] << host
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_services( xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -145,9 +152,11 @@ public
|
|||
ret[:services] << service
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_vulns(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -177,6 +186,7 @@ public
|
|||
ret[:vulns] << vuln
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_workspaces
|
||||
|
@ -215,6 +225,7 @@ public
|
|||
end
|
||||
|
||||
def rpc_set_workspace(wspace)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
workspace = self.framework.db.find_workspace(wspace)
|
||||
if(workspace)
|
||||
|
@ -222,9 +233,11 @@ public
|
|||
return { 'result' => "success" }
|
||||
end
|
||||
{ 'result' => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_del_workspace(wspace)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
# Delete workspace
|
||||
workspace = self.framework.db.find_workspace(wspace)
|
||||
|
@ -240,16 +253,20 @@ public
|
|||
workspace.destroy
|
||||
end
|
||||
{ 'result' => "success" }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_add_workspace(wspace)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
wspace = self.framework.db.add_workspace(wspace)
|
||||
return { 'result' => 'success' } if(wspace)
|
||||
{ 'result' => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_host(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
|
||||
ret = {}
|
||||
|
@ -273,25 +290,30 @@ public
|
|||
ret[:host] << host
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_report_host(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
|
||||
res = self.framework.db.report_host(opts)
|
||||
return { :result => 'success' } if(res)
|
||||
{ :result => 'failed' }
|
||||
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_report_service(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
res = self.framework.db.report_service(opts)
|
||||
return { :result => 'success' } if(res)
|
||||
{ :result => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_service(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
|
||||
ret = {}
|
||||
|
@ -333,9 +355,11 @@ public
|
|||
ret[:service] << service
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_note(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
|
||||
ret = {}
|
||||
|
@ -383,9 +407,11 @@ public
|
|||
ret[:note] << note
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_client(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
ret = {}
|
||||
ret[:client] = []
|
||||
|
@ -402,17 +428,21 @@ public
|
|||
ret[:client] << client
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_report_client(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
res = self.framework.db.report_client(opts)
|
||||
return { :result => 'success' } if(res)
|
||||
{ :result => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
#DOC NOTE: :data and :ntype are REQUIRED
|
||||
def rpc_report_note(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
if (opts[:host] or opts[:address]) and opts[:port] and opts[:proto]
|
||||
addr = opts[:host] || opts[:address]
|
||||
|
@ -425,9 +455,11 @@ public
|
|||
res = self.framework.db.report_note(opts)
|
||||
return { :result => 'success' } if(res)
|
||||
{ :result => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_notes(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -454,19 +486,24 @@ public
|
|||
ret[:notes] << note
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_report_auth_info(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
res = self.framework.db.report_auth_info(opts)
|
||||
return { :result => 'success' } if(res)
|
||||
{ :result => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_auth_info(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
ret = {}
|
||||
ret[:auth_info] = []
|
||||
# XXX: This method doesn't exist...
|
||||
ai = self.framework.db.get_auth_info(opts)
|
||||
ai.each do |i|
|
||||
info = {}
|
||||
|
@ -476,14 +513,18 @@ public
|
|||
ret[:auth_info] << info
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_ref(name)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
self.framework.db.get_ref(name)
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_del_vuln(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
hosts = []
|
||||
services = []
|
||||
|
@ -548,9 +589,11 @@ public
|
|||
end
|
||||
|
||||
return { :result => 'success', :deleted => deleted }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_del_note(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
hosts = []
|
||||
services = []
|
||||
|
@ -614,9 +657,11 @@ public
|
|||
end
|
||||
|
||||
return { :result => 'success', :deleted => deleted }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_del_service(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
hosts = []
|
||||
services = []
|
||||
|
@ -668,9 +713,11 @@ public
|
|||
end
|
||||
|
||||
return { :result => 'success', :deleted => deleted }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_del_host(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
opts = fix_options(xopts)
|
||||
wspace = find_workspace(opts[:workspace])
|
||||
|
@ -695,18 +742,22 @@ public
|
|||
end
|
||||
|
||||
return { :result => 'success', :deleted => deleted }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_report_vuln(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
opts = fix_options(xopts)
|
||||
opts[:workspace] = find_workspace(opts[:workspace]) if opts[:workspace]
|
||||
res = self.framework.db.report_vuln(opts)
|
||||
return { :result => 'success' } if(res)
|
||||
{ :result => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_events(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -726,17 +777,21 @@ public
|
|||
ret[:events] << event
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_report_event(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
res = self.framework.db.report_event(opts)
|
||||
{ :result => 'success' } if(res)
|
||||
}
|
||||
end
|
||||
|
||||
#NOTE Path is required
|
||||
#NOTE To match a service need host, port, proto
|
||||
def rpc_report_loot(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
if opts[:host] && opts[:port] && opts[:proto]
|
||||
opts[:service] = self.framework.db.find_or_create_service(opts)
|
||||
|
@ -744,9 +799,11 @@ public
|
|||
|
||||
res = self.framework.db.report_loot(opts)
|
||||
{ :result => 'success' } if(res)
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_loots(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -767,18 +824,22 @@ public
|
|||
ret[:loots] << loot
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
# requires host, port, user, pass, ptype, and active
|
||||
def rpc_report_cred(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
res = framework.db.find_or_create_cred(opts)
|
||||
return { :result => 'success' } if res
|
||||
{ :result => 'failed' }
|
||||
}
|
||||
end
|
||||
|
||||
#right now workspace is the only option supported
|
||||
def rpc_creds(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -800,15 +861,19 @@ public
|
|||
ret[:creds] << cred
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_import_data(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
self.framework.db.import(opts)
|
||||
return { :result => 'success' }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_get_vuln(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
|
||||
ret = {}
|
||||
|
@ -854,9 +919,11 @@ public
|
|||
ret[:vuln] << vuln
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_clients(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
opts, wspace = init_db_opts_workspace(xopts)
|
||||
limit = opts.delete(:limit) || 100
|
||||
offset = opts.delete(:offset) || 0
|
||||
|
@ -881,9 +948,11 @@ public
|
|||
ret[:clients] << client
|
||||
end
|
||||
ret
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_del_client(xopts)
|
||||
::ActiveRecord::Base.connection_pool.with_connection {
|
||||
db_check
|
||||
opts = fix_options(xopts)
|
||||
wspace = find_workspace(opts[:workspace])
|
||||
|
@ -921,6 +990,7 @@ public
|
|||
end
|
||||
|
||||
{ :result => 'success', :deleted => deleted }
|
||||
}
|
||||
end
|
||||
|
||||
def rpc_driver(xopts)
|
||||
|
@ -975,12 +1045,14 @@ public
|
|||
end
|
||||
|
||||
cdb = ""
|
||||
if ActiveRecord::Base.connected? and ActiveRecord::Base.connection.active?
|
||||
if ActiveRecord::Base.connection.respond_to? :current_database
|
||||
cdb = ActiveRecord::Base.connection.current_database
|
||||
else
|
||||
cdb = ActiveRecord::Base.connection.instance_variable_get(:@config)[:database]
|
||||
end
|
||||
if ::ActiveRecord::Base.connected?
|
||||
::ActiveRecord::Base.connection_pool.with_connection { |conn|
|
||||
if conn.respond_to? :current_database
|
||||
cdb = conn.current_database
|
||||
else
|
||||
cdb = conn.instance_variable_get(:@config)[:database]
|
||||
end
|
||||
}
|
||||
return {:driver => self.framework.db.driver.to_s , :db => cdb }
|
||||
else
|
||||
return {:driver => self.framework.db.driver.to_s}
|
||||
|
|
Loading…
Reference in New Issue