Skip to content

Commit 1d430db

Browse files
committed
Run migrations when connection already established in console
MSP-10955 `Msf::Ui::Console::Driver#initialize` doesn't call `framework.db.connect` if it can't find the the `database.yml`, but when using `msfpro`, the connection is already established, so the console doesn't need to know where the database file is and should just run the migrations so that `framework.db.migrate` can be set and `framework.db.active` will return `true`.
1 parent 5af4959 commit 1d430db

File tree

2 files changed

+74
-48
lines changed

2 files changed

+74
-48
lines changed

lib/msf/core/db_manager.rb

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -228,28 +228,44 @@ def connect(opts={})
228228
# Configure the database adapter
229229
ActiveRecord::Base.establish_connection(nopts)
230230
end
231-
232-
# Migrate the database, if needed
233-
migrate
234-
235-
# Set the default workspace
236-
framework.db.workspace = framework.db.default_workspace
237-
238-
# Flag that migration has completed
239-
self.migrated = true
240231
rescue ::Exception => e
241232
self.error = e
242233
elog("DB.connect threw an exception: #{e}")
243234
dlog("Call stack: #{$@.join"\n"}", LEV_1)
244235
return false
245236
ensure
237+
after_establish_connection
238+
246239
# Database drivers can reset our KCODE, do not let them
247240
$KCODE = 'NONE' if RUBY_VERSION =~ /^1\.8\./
248241
end
249242

250243
true
251244
end
252245

246+
# Finishes {#connect} after `ActiveRecord::Base.establish_connection` has succeeded by {#migrate migrating database}
247+
# and setting {#workspace}.
248+
#
249+
# @return [void]
250+
def after_establish_connection
251+
self.migrated = false
252+
253+
begin
254+
# Migrate the database, if needed
255+
migrate
256+
257+
# Set the default workspace
258+
framework.db.workspace = framework.db.default_workspace
259+
rescue ::Exception => exception
260+
self.error = exception
261+
elog("DB.connect threw an exception: #{exception}")
262+
dlog("Call stack: #{exception.backtrace.join("\n")}", LEV_1)
263+
else
264+
# Flag that migration has completed
265+
self.migrated = true
266+
end
267+
end
268+
253269
#
254270
# Attempt to create the database
255271
#

lib/msf/ui/console/driver.rb

Lines changed: 49 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -176,50 +176,60 @@ def initialize(prompt = DefaultPrompt, prompt_char = DefaultPromptChar, opts = {
176176
end
177177
end
178178

179-
# Look for our database configuration in the following places, in order:
180-
# command line arguments
181-
# environment variable
182-
# configuration directory (usually ~/.msf3)
183-
dbfile = opts['DatabaseYAML']
184-
dbfile ||= ENV["MSF_DATABASE_CONFIG"]
185-
dbfile ||= File.join(Msf::Config.get_config_root, "database.yml")
186-
if (dbfile and File.exists? dbfile)
187-
if File.readable?(dbfile)
188-
dbinfo = YAML.load(File.read(dbfile))
189-
dbenv = opts['DatabaseEnv'] || Rails.env
190-
db = dbinfo[dbenv]
191-
else
192-
print_error("Warning, #{dbfile} is not readable. Try running as root or chmod.")
193-
end
194-
if not db
195-
print_error("No database definition for environment #{dbenv}")
196-
else
197-
if not framework.db.connect(db)
198-
if framework.db.error.to_s =~ /RubyGem version.*pg.*0\.11/i
199-
print_error("***")
200-
print_error("*")
201-
print_error("* Metasploit now requires version 0.11 or higher of the 'pg' gem for database support")
202-
print_error("* There a three ways to accomplish this upgrade:")
203-
print_error("* 1. If you run Metasploit with your system ruby, simply upgrade the gem:")
204-
print_error("* $ rvmsudo gem install pg ")
205-
print_error("* 2. Use the Community Edition web interface to apply a Software Update")
206-
print_error("* 3. Uninstall, download the latest version, and reinstall Metasploit")
207-
print_error("*")
208-
print_error("***")
209-
print_error("")
210-
print_error("")
211-
end
212-
213-
print_error("Failed to connect to the database: #{framework.db.error}")
179+
if framework.db.connection_established?
180+
framework.db.after_establish_connection
181+
else
182+
# Look for our database configuration in the following places, in order:
183+
# command line arguments
184+
# environment variable
185+
# configuration directory (usually ~/.msf3)
186+
dbfile = opts['DatabaseYAML']
187+
dbfile ||= ENV["MSF_DATABASE_CONFIG"]
188+
dbfile ||= File.join(Msf::Config.get_config_root, "database.yml")
189+
190+
if (dbfile and File.exists? dbfile)
191+
if File.readable?(dbfile)
192+
dbinfo = YAML.load(File.read(dbfile))
193+
dbenv = opts['DatabaseEnv'] || Rails.env
194+
db = dbinfo[dbenv]
214195
else
215-
self.framework.modules.refresh_cache_from_database
196+
print_error("Warning, #{dbfile} is not readable. Try running as root or chmod.")
197+
end
216198

217-
if self.framework.modules.cache_empty?
218-
print_status("The initial module cache will be built in the background, this can take 2-5 minutes...")
219-
end
199+
if not db
200+
print_error("No database definition for environment #{dbenv}")
201+
else
202+
framework.db.connect(db)
220203
end
221204
end
222205
end
206+
207+
# framework.db.active will be true if after_establish_connection ran directly when connection_established? was
208+
# already true or if framework.db.connect called after_establish_connection.
209+
if framework.db.active
210+
self.framework.modules.refresh_cache_from_database
211+
212+
if self.framework.modules.cache_empty?
213+
print_status("The initial module cache will be built in the background, this can take 2-5 minutes...")
214+
end
215+
elsif !framework.db.error.nil?
216+
if framework.db.error.to_s =~ /RubyGem version.*pg.*0\.11/i
217+
print_error("***")
218+
print_error("*")
219+
print_error("* Metasploit now requires version 0.11 or higher of the 'pg' gem for database support")
220+
print_error("* There a three ways to accomplish this upgrade:")
221+
print_error("* 1. If you run Metasploit with your system ruby, simply upgrade the gem:")
222+
print_error("* $ rvmsudo gem install pg ")
223+
print_error("* 2. Use the Community Edition web interface to apply a Software Update")
224+
print_error("* 3. Uninstall, download the latest version, and reinstall Metasploit")
225+
print_error("*")
226+
print_error("***")
227+
print_error("")
228+
print_error("")
229+
end
230+
231+
print_error("Failed to connect to the database: #{framework.db.error}")
232+
end
223233
end
224234

225235
# Initialize the module paths only if we didn't get passed a Framework instance

0 commit comments

Comments
 (0)