diff --git a/lib/msf/core/db_manager.rb b/lib/msf/core/db_manager.rb index fc9ca56424..a618613a64 100644 --- a/lib/msf/core/db_manager.rb +++ b/lib/msf/core/db_manager.rb @@ -3,6 +3,7 @@ require 'msf/base/config' require 'msf/core' require 'msf/core/db' +require 'msf/core/db_manager/migration' require 'msf/core/task_manager' require 'fileutils' require 'shellwords' @@ -17,6 +18,9 @@ module Msf ### class DBManager + # Provides :framework and other accessors + include Msf::DBManager::Migration + include Msf::Framework::Offspring # Mainly, it's Ruby 1.9.1 that cause a lot of problems now, along with Ruby 1.8.6. # Ruby 1.8.7 actually seems okay, but why tempt fate? Let's say 1.9.3 and beyond. @@ -28,9 +32,6 @@ class DBManager end end - # Provides :framework and other accessors - include Framework::Offspring - # Returns true if we are ready to load/store data def active return false if not @usable @@ -53,9 +54,6 @@ class DBManager # Stores a TaskManager for serializing database events attr_accessor :sink - # Flag to indicate database migration has completed - attr_accessor :migrated - # Flag to indicate that modules are cached attr_accessor :modules_cached @@ -287,33 +285,6 @@ class DBManager end end - # Migrate database to latest schema version. - # - # @param verbose [Boolean] see ActiveRecord::Migration.verbose - # @return [Array error - self.error = error - elog("DB.migrate threw an exception: #{error}") - dlog("Call stack:\n#{error.backtrace.join "\n"}") - end - end - - return ran - end - def workspace=(workspace) @workspace_name = workspace.name end diff --git a/lib/msf/core/db_manager/migration.rb b/lib/msf/core/db_manager/migration.rb new file mode 100644 index 0000000000..517c379638 --- /dev/null +++ b/lib/msf/core/db_manager/migration.rb @@ -0,0 +1,58 @@ +module Msf + class DBManager + module Migration + # Migrate database to latest schema version. + # + # @param verbose [Boolean] see ActiveRecord::Migration.verbose + # @return [Array error + self.error = error + elog("DB.migrate threw an exception: #{error}") + dlog("Call stack:\n#{error.backtrace.join "\n"}") + end + end + + # Since the connections that existed before the migrations ran could + # have outdated column information, reset column information for all + # ActiveRecord::Base descendents to prevent missing method errors for + # column methods for columns created in migrations after the column + # information was cached. + reset_column_information + + return ran + end + + # Flag to indicate database migration has completed + # + # @return [Boolean] + attr_accessor :migrated + + private + + # Resets the column information for all descendants of ActiveRecord::Base + # since some of the migrations may have cached column information that + # has been updated by later migrations. + # + # @return [void] + def reset_column_information + ActiveRecord::Base.descendants.each do |descendant| + descendant.reset_column_information + end + end + end + end +end \ No newline at end of file diff --git a/spec/lib/msf/db_manager_spec.rb b/spec/lib/msf/db_manager_spec.rb index 914b8feae8..2f0472b9a1 100644 --- a/spec/lib/msf/db_manager_spec.rb +++ b/spec/lib/msf/db_manager_spec.rb @@ -18,6 +18,7 @@ describe Msf::DBManager do db_manager end + it_should_behave_like 'Msf::DBManager::Migration' it_should_behave_like 'Msf::DBManager::ImportMsfXml' context '#initialize_metasploit_data_models' do diff --git a/spec/support/shared/examples/msf/db_manager/migration.rb b/spec/support/shared/examples/msf/db_manager/migration.rb new file mode 100644 index 0000000000..6cf698dace --- /dev/null +++ b/spec/support/shared/examples/msf/db_manager/migration.rb @@ -0,0 +1,143 @@ +shared_examples_for 'Msf::DBManager::Migration' do + it { should be_a Msf::DBManager::Migration } + + context '#migrate' do + def migrate + db_manager.migrate + end + + it 'should create a connection' do + ActiveRecord::Base.connection_pool.should_receive(:with_connection).twice + + migrate + end + + it 'should call ActiveRecord::Migrator.migrate' do + ActiveRecord::Migrator.should_receive(:migrate).with( + ActiveRecord::Migrator.migrations_paths + ) + + migrate + end + + it 'should return migrations that were ran from ActiveRecord::Migrator.migrate' do + migrations = [mock('Migration 1')] + ActiveRecord::Migrator.stub(:migrate => migrations) + + migrate.should == migrations + end + + it 'should reset the column information' do + db_manager.should_receive(:reset_column_information) + + migrate + end + + context 'with StandardError from ActiveRecord::Migration.migrate' do + let(:error) do + StandardError.new(message) + end + + let(:message) do + "Error during migration" + end + + before(:each) do + ActiveRecord::Migrator.stub(:migrate).and_raise(error) + end + + it 'should set Msf::DBManager#error' do + migrate + + db_manager.error.should == error + end + + it 'should log error message at error level' do + db_manager.should_receive(:elog) do |error_message| + error_message.should include(error.to_s) + end + + migrate + end + + it 'should log error backtrace at debug level' do + db_manager.should_receive(:dlog) do |debug_message| + debug_message.should include('Call stack') + end + + migrate + end + end + + context 'with verbose' do + def migrate + db_manager.migrate(verbose) + end + + context 'false' do + let(:verbose) do + false + end + + it 'should set ActiveRecord::Migration.verbose to false' do + ActiveRecord::Migration.should_receive(:verbose=).with(verbose) + + migrate + end + end + + context 'true' do + let(:verbose) do + true + end + + it 'should set ActiveRecord::Migration.verbose to true' do + ActiveRecord::Migration.should_receive(:verbose=).with(verbose) + + migrate + end + end + end + + context 'without verbose' do + it 'should set ActiveRecord::Migration.verbose to false' do + ActiveRecord::Migration.should_receive(:verbose=).with(false) + + db_manager.migrate + end + end + end + + context '#migrated' do + it { should respond_to :migrated } + it { should respond_to :migrated= } + end + + context '#reset_column_information' do + def reset_column_information + db_manager.send(:reset_column_information) + end + + it 'should use ActiveRecord::Base.descendants to find both direct and indirect subclasses' do + ActiveRecord::Base.should_receive(:descendants).and_return([]) + + reset_column_information + end + + it 'should reset column information on each descendant of ActiveRecord::Base' do + descendants = [] + + 1.upto(2) do |i| + descendants << mock("Descendant #{i}") + end + + ActiveRecord::Base.stub(:descendants => descendants) + + descendants.each do |descendant| + descendant.should_receive(:reset_column_information) + end + + reset_column_information + end + end +end \ No newline at end of file