Msf::ModuleManager::Cache spec coverage

[#47979793]

Author: limhoff-r7

Gemfile

@@ -42,6 +42,9 @@ group :test do
 	gem 'database_cleaner'
   # testing framework
   gem 'rspec', '>= 2.12'
+  # add matchers from shoulda, such as query_the_database, which is useful for
+  # testing that the Msf::DBManager activation is respected.
+  gem 'shoulda-matchers'
   # code coverage for tests
   # any version newer than 0.5.4 gives an Encoding error when trying to read the source files.
   gem 'simplecov', '0.5.4', :require => false

Gemfile.lock

@@ -13,17 +13,22 @@ GEM
       i18n (= 0.6.1)
       multi_json (~> 1.0)
     arel (3.0.2)
+    bourne (1.4.0)
+      mocha (~> 0.13.2)
     builder (3.0.4)
     database_cleaner (0.9.1)
     diff-lcs (1.2.2)
     factory_girl (4.2.0)
       activesupport (>= 3.0.0)
     i18n (0.6.1)
     json (1.7.7)
+    metaclass (0.0.1)
     metasploit_data_models (0.6.14)
       activerecord (>= 3.2.13)
       activesupport
       pg
+    mocha (0.13.3)
+      metaclass (~> 0.0.1)
     msgpack (0.5.4)
     multi_json (1.0.4)
     nokogiri (1.5.9)
@@ -40,6 +45,9 @@ GEM
     rspec-expectations (2.13.0)
       diff-lcs (>= 1.1.3, < 2.0)
     rspec-mocks (2.13.0)
+    shoulda-matchers (1.5.2)
+      activesupport (>= 3.0.0)
+      bourne (~> 1.3)
     simplecov (0.5.4)
       multi_json (~> 1.0.3)
       simplecov-html (~> 0.5.3)
@@ -65,5 +73,6 @@ DEPENDENCIES
   redcarpet
   robots
   rspec (>= 2.12)
+  shoulda-matchers
   simplecov (= 0.5.4)
   yard

lib/msf/core/module_manager/cache.rb

@@ -46,13 +46,19 @@ def load_cached_module(type, reference_name)
     loaded
   end
 
-  # Rebuild the cache for the module set
+  # @overload refresh_cache_from_module_files
+  #   Rebuilds database and in-memory cache for all modules.
   #
-  # @return [void]
-  def refresh_cache_from_module_files(mod = nil)
+  #   @return [void]
+  # @overload refresh_cache_from_module_files(module_class_or_instance)
+  #   Rebuilds database and in-memory cache for given module_class_or_instance.
+  #
+  #   @param (see Msf::DBManager#update_module_details)
+  #   @return [void]
+  def refresh_cache_from_module_files(module_class_or_instance = nil)
     if framework_migrated?
-      if mod
-        framework.db.update_module_details(mod)
+      if module_class_or_instance
+        framework.db.update_module_details(module_class_or_instance)
       else
         framework.db.update_all_module_details
       end
@@ -61,7 +67,7 @@ def refresh_cache_from_module_files(mod = nil)
     end
   end
 
-  # Reset the module cache
+  # Refreshes the in-memory cache from the database cache.
   #
   # @return [void]
   def refresh_cache_from_database
@@ -97,32 +103,38 @@ def module_info_by_path_from_database!
     self.module_info_by_path = {}
 
     if framework_migrated?
-      # TODO record module parent_path in {Mdm::ModuleDetail} so it does not need to be derived from file.
-      ::Mdm::ModuleDetail.find(:all).each do |module_detail|
-        path = module_detail.file
-        type = module_detail.mtype
-        reference_name = module_detail.refname
-
-        typed_path = Msf::Modules::Loader::Base.typed_path(type, reference_name)
-        escaped_typed_path = Regexp.escape(typed_path)
-        parent_path = path.gsub(/#{escaped_typed_path}$/, '')
-
-        module_info_by_path[path] = {
-            :reference_name => reference_name,
-            :type => type,
-            :parent_path => parent_path,
-            :modification_time => module_detail.mtime
-        }
-
-        typed_module_set = module_set(type)
-
-        # Don't want to trigger as {Msf::ModuleSet#create} so check for
-        # key instead of using ||= which would call {Msf::ModuleSet#[]}
-        # which would potentially call {Msf::ModuleSet#create}.
-        unless typed_module_set.has_key? reference_name
-          typed_module_set[reference_name] = Msf::SymbolicModule
-        end
-      end
+	    ActiveRecord::Base.connection_pool.with_connection do
+		    # TODO record module parent_path in {Mdm::ModuleDetail} so it does not need to be derived from file.
+		    # Use find_each so Mdm::ModuleDetails are returned in batches, which will
+		    # handle the growing number of modules better than all.each.
+		    Mdm::ModuleDetail.find_each do |module_detail|
+			    path = module_detail.file
+			    type = module_detail.mtype
+			    reference_name = module_detail.refname
+
+			    typed_path = Msf::Modules::Loader::Base.typed_path(type, reference_name)
+			    # join to '' so that typed_path_prefix starts with file separator
+			    typed_path_suffix = File.join('', typed_path)
+			    escaped_typed_path = Regexp.escape(typed_path_suffix)
+			    parent_path = path.gsub(/#{escaped_typed_path}$/, '')
+
+			    module_info_by_path[path] = {
+					    :reference_name => reference_name,
+					    :type => type,
+					    :parent_path => parent_path,
+					    :modification_time => module_detail.mtime
+			    }
+
+			    typed_module_set = module_set(type)
+
+			    # Don't want to trigger as {Msf::ModuleSet#create} so check for
+			    # key instead of using ||= which would call {Msf::ModuleSet#[]}
+			    # which would potentially call {Msf::ModuleSet#create}.
+			    unless typed_module_set.has_key? reference_name
+				    typed_module_set[reference_name] = Msf::SymbolicModule
+			    end
+		    end
+	    end
     end
 
     self.module_info_by_path

spec/lib/msf/core/module_manager_spec.rb

@@ -16,6 +16,8 @@
 require 'msf/core'
 
 describe Msf::ModuleManager do
+	include_context 'Msf::Simple::Framework'
+
   let(:archive_basename) do
     [basename_prefix, archive_extension]
   end
@@ -28,14 +30,11 @@
     'rspec'
   end
 
-  let(:framework) do
-    Msf::Framework.new
-  end
-
   subject(:module_manager) do
-    described_class.new(framework)
+    framework.modules
   end
 
+  it_should_behave_like 'Msf::ModuleManager::Cache'
   it_should_behave_like 'Msf::ModuleManager::Loading'
   it_should_behave_like 'Msf::ModuleManager::ModulePaths'
 end

spec/support/shared/contexts/database_cleaner.rb

@@ -1,3 +1,5 @@
+require 'metasploit/framework/database'
+
 shared_context 'DatabaseCleaner' do
 	def with_established_connection
 		begin