Add with_connection wrappers to the methods I missed

unstable
James Lee 2012-04-19 15:45:17 -06:00
parent 876c59b192
commit 06b3ed2e13
1 changed files with 35 additions and 4 deletions

View File

@ -506,7 +506,7 @@ class DBManager
# #
def report_service(opts) def report_service(opts)
return if not active return if not active
#::ActiveRecord::Base.connection_pool.with_connection { |conn| ::ActiveRecord::Base.connection_pool.with_connection { |conn|
addr = opts.delete(:host) || return addr = opts.delete(:host) || return
hname = opts.delete(:host_name) hname = opts.delete(:host_name)
hmac = opts.delete(:mac) hmac = opts.delete(:mac)
@ -555,14 +555,15 @@ class DBManager
dlog("Unknown attribute for Service: #{k}") dlog("Unknown attribute for Service: #{k}")
end end
} }
service.state = ServiceState::Open if service.state.nil? service.state ||= ServiceState::Open
service.info = "" if service.info.nil? service.info ||= ""
if (service and service.changed?) if (service and service.changed?)
msf_import_timestamps(opts,service) msf_import_timestamps(opts,service)
service.save! service.save!
end end
ret[:service] = service ret[:service] = service
#} }
end end
def get_service(wspace, host, proto, port) def get_service(wspace, host, proto, port)
@ -1805,6 +1806,7 @@ class DBManager
def report_web_page(opts) def report_web_page(opts)
return if not active return if not active
::ActiveRecord::Base.connection_pool.with_connection {
wspace = opts.delete(:workspace) || workspace wspace = opts.delete(:workspace) || workspace
path = opts[:path] path = opts[:path]
@ -1846,6 +1848,7 @@ class DBManager
page.save! page.save!
ret[:web_page] = page ret[:web_page] = page
}
end end
@ -1871,6 +1874,7 @@ class DBManager
def report_web_form(opts) def report_web_form(opts)
return if not active return if not active
::ActiveRecord::Base.connection_pool.with_connection {
wspace = opts.delete(:workspace) || workspace wspace = opts.delete(:workspace) || workspace
path = opts[:path] path = opts[:path]
@ -1924,6 +1928,7 @@ class DBManager
msf_import_timestamps(opts, form) msf_import_timestamps(opts, form)
form.save! form.save!
ret[:web_form] = form ret[:web_form] = form
}
end end
@ -1953,6 +1958,7 @@ class DBManager
def report_web_vuln(opts) def report_web_vuln(opts)
return if not active return if not active
::ActiveRecord::Base.connection_pool.with_connection {
wspace = opts.delete(:workspace) || workspace wspace = opts.delete(:workspace) || workspace
path = opts[:path] path = opts[:path]
@ -2024,6 +2030,7 @@ class DBManager
vuln.save! vuln.save!
ret[:web_vuln] = vuln ret[:web_vuln] = vuln
}
end end
# #
@ -2031,12 +2038,14 @@ class DBManager
# Selected host # Selected host
# #
def selected_host def selected_host
::ActiveRecord::Base.connection_pool.with_connection {
selhost = ::Mdm::WmapTarget.where("selected != 0").first() selhost = ::Mdm::WmapTarget.where("selected != 0").first()
if selhost if selhost
return selhost.host return selhost.host
else else
return return
end end
}
end end
# #
@ -2044,7 +2053,9 @@ class DBManager
# Selected target # Selected target
# #
def selected_wmap_target def selected_wmap_target
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapTarget.find.where("selected != 0") ::Mdm::WmapTarget.find.where("selected != 0")
}
end end
# #
@ -2088,7 +2099,9 @@ class DBManager
# This method wiil be remove on second phase of db merging. # This method wiil be remove on second phase of db merging.
# #
def request_distinct_targets def request_distinct_targets
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapRequest.select('DISTINCT host,address,port,ssl') ::Mdm::WmapRequest.select('DISTINCT host,address,port,ssl')
}
end end
# #
@ -2146,7 +2159,9 @@ class DBManager
# This method returns a list of all requests from target # This method returns a list of all requests from target
# #
def target_requests(extra_condition) def target_requests(extra_condition)
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapRequest.where("wmap_requests.host = ? AND wmap_requests.port = ? #{extra_condition}",selected_host,selected_port) ::Mdm::WmapRequest.where("wmap_requests.host = ? AND wmap_requests.port = ? #{extra_condition}",selected_host,selected_port)
}
end end
# #
@ -2165,7 +2180,9 @@ class DBManager
# This method allows to query directly the requests table. To be used mainly by modules # This method allows to query directly the requests table. To be used mainly by modules
# #
def request_sql(host,port,extra_condition) def request_sql(host,port,extra_condition)
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapRequest.where("wmap_requests.host = ? AND wmap_requests.port = ? #{extra_condition}", host , port) ::Mdm::WmapRequest.where("wmap_requests.host = ? AND wmap_requests.port = ? #{extra_condition}", host , port)
}
end end
# #
@ -2173,7 +2190,9 @@ class DBManager
# This methods returns a list of all targets in the database # This methods returns a list of all targets in the database
# #
def requests def requests
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapRequest.find(:all) ::Mdm::WmapRequest.find(:all)
}
end end
# #
@ -2192,7 +2211,9 @@ class DBManager
# This methods returns a list of all targets in the database # This methods returns a list of all targets in the database
# #
def targets def targets
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapTarget.find(:all) ::Mdm::WmapTarget.find(:all)
}
end end
# #
@ -2200,7 +2221,9 @@ class DBManager
# This methods deletes all targets from targets table in the database # This methods deletes all targets from targets table in the database
# #
def delete_all_targets def delete_all_targets
::ActiveRecord::Base.connection_pool.with_connection {
::Mdm::WmapTarget.delete_all ::Mdm::WmapTarget.delete_all
}
end end
# #
@ -2208,8 +2231,10 @@ class DBManager
# Find a target matching this id # Find a target matching this id
# #
def get_target(id) def get_target(id)
::ActiveRecord::Base.connection_pool.with_connection {
target = ::Mdm::WmapTarget.where("id = ?", id).first() target = ::Mdm::WmapTarget.where("id = ?", id).first()
return target return target
}
end end
# #
@ -2217,6 +2242,7 @@ class DBManager
# Create a target # Create a target
# #
def create_target(host,port,ssl,sel) def create_target(host,port,ssl,sel)
::ActiveRecord::Base.connection_pool.with_connection {
tar = ::Mdm::WmapTarget.create( tar = ::Mdm::WmapTarget.create(
:host => host, :host => host,
:address => host, :address => host,
@ -2225,6 +2251,7 @@ class DBManager
:selected => sel :selected => sel
) )
#framework.events.on_db_target(rec) #framework.events.on_db_target(rec)
}
end end
@ -2233,6 +2260,7 @@ class DBManager
# Create a request (by hand) # Create a request (by hand)
# #
def create_request(host,port,ssl,meth,path,headers,query,body,respcode,resphead,response) def create_request(host,port,ssl,meth,path,headers,query,body,respcode,resphead,response)
::ActiveRecord::Base.connection_pool.with_connection {
req = ::Mdm::WmapRequest.create( req = ::Mdm::WmapRequest.create(
:host => host, :host => host,
:address => host, :address => host,
@ -2248,6 +2276,7 @@ class DBManager
:response => response :response => response
) )
#framework.events.on_db_request(rec) #framework.events.on_db_request(rec)
}
end end
# #
@ -2255,7 +2284,9 @@ class DBManager
# Quick way to query the database (used by wmap_sql) # Quick way to query the database (used by wmap_sql)
# #
def sql_query(sqlquery) def sql_query(sqlquery)
::ActiveRecord::Base.connection_pool.with_connection {
ActiveRecord::Base.connection.select_all(sqlquery) ActiveRecord::Base.connection.select_all(sqlquery)
}
end end