diff --git a/lib/msf/core/db_manager.rb b/lib/msf/core/db_manager.rb index e76cab0c6c..701f0631c0 100644 --- a/lib/msf/core/db_manager.rb +++ b/lib/msf/core/db_manager.rb @@ -1,6 +1,7 @@ require 'msf/core' require 'msf/core/db' require 'msf/core/task_manager' +require 'fileutils' module Msf @@ -39,11 +40,16 @@ class DBManager # Flag to indicate database migration has completed attr_accessor :migrated + + # Array of additional migration paths + attr_accessor :migration_paths def initialize(framework, opts = {}) self.framework = framework self.migrated = false + self.migration_paths = [ ::File.join(Msf::Config.install_root, "data", "sql", "migrate") ] + @usable = false # Don't load the database if the user said they didn't need it. @@ -225,16 +231,37 @@ class DBManager # Migrate database to latest schema version # def migrate(verbose=false) + + temp_dir = ::File.expand_path(::File.join( Msf::Config.config_directory, "schema", "#{Time.now.to_i}_#{$$}" )) + ::FileUtils.rm_rf(temp_dir) + ::FileUtils.mkdir_p(temp_dir) + + self.migration_paths.each do |mpath| + dir = Dir.new(mpath) rescue nil + if not dir + elog("Could not access migration path #{mpath}") + next + end + + dir.entries.each do |ent| + next unless ent =~ /^\d+.*\.rb$/ + ::FileUtils.cp( ::File.join(mpath, ent), ::File.join(temp_dir, ent) ) + end + end + + success = true begin - migrate_dir = ::File.join(Msf::Config.install_root, "data", "sql", "migrate") ActiveRecord::Migration.verbose = verbose - ActiveRecord::Migrator.migrate(migrate_dir, nil) + ActiveRecord::Migrator.migrate(temp_dir, nil) rescue ::Exception => e self.error = e elog("DB.migrate threw an exception: #{e}") dlog("Call stack:\n#{e.backtrace.join "\n"}") - return false + success = false end + + ::FileUtils.rm_rf(temp_dir) + return true end diff --git a/lib/msf/ui/console/driver.rb b/lib/msf/ui/console/driver.rb index 7c4a8e298d..9576101021 100644 --- a/lib/msf/ui/console/driver.rb +++ b/lib/msf/ui/console/driver.rb @@ -174,6 +174,12 @@ class Driver < Msf::Ui::Driver # Parse any specified database.yml file if framework.db.usable and not opts['SkipDatabaseInit'] + + # Append any migration paths necessary to bring the database online + if opts['DatabaseMigrationPaths'] + opts['DatabaseMigrationPaths'].each {|m| framework.db.migrations_paths << m } + end + # Look for our database configuration in the following places, in order: # command line arguments # environment variable diff --git a/msfconsole b/msfconsole index 220ca9b5e5..f3cd7394eb 100755 --- a/msfconsole +++ b/msfconsole @@ -71,6 +71,11 @@ class OptsConsole options['DatabaseYAML'] = m end + opts.on("-M", "--migration-path ", "Specify a directory containing additional DB migrations") do |m| + options['DatabaseMigrationPaths'] ||= [] + options['DatabaseMigrationPaths'] << m + end + opts.on("-e", "--environment ", "Specify the database environment to load from the YAML") do |m| options['DatabaseEnv'] = m end