Merge pull request #20372 from cgranleese-r7/add-cache-validation-logic

Adds cache invalidation logic

Author: smcintyre-r7

lib/msf/core/modules/metadata/store.rb

@@ -1,4 +1,6 @@
 require 'json'
+require 'parallel'
+require 'zlib'
 
 #
 # Handles storage of module metadata on disk. A base metadata file is always included - this was added to ensure a much
@@ -14,6 +16,7 @@ def initialize
 
   BaseMetaDataFile = 'modules_metadata_base.json'
   UserMetaDataFile = 'modules_metadata.json'
+  CacheMetaDataFile = 'module_metadata_cache.json'
 
   #
   # Initializes from user store (under ~/store/.msf4) if it exists. else base file (under $INSTALL_ROOT/db) is copied and loaded.
@@ -124,4 +127,164 @@ def load_cache_from_file_store
     }
   end
 
+  # This method checks if the current module and library files match the cached checksum.
+  # It uses a per-file CRC32 cache to avoid recalculating checksums for files that haven't changed.
+  # If no cache exists, it will create one in the user's directory.
+  #
+  # @return [Boolean] True if the current checksum matches the cached one
+  def self.valid_checksum?
+    current_checksum = get_current_checksum
+    cached_sha = get_cached_checksum
+
+    # If no cached checksum exists, create the cache file with current checksum
+    if cached_sha.nil?
+      update_cache_checksum(current_checksum)
+      return false
+    end
+
+    checksums_match?(current_checksum, cached_sha)
+  end
+
+  # Calculate the current checksum for all module and library files
+  # This calculates checksums for each file and generates an overall checksum
+  # from the individual file checksums. Does NOT update the cached checksum.
+  #
+  # @return [Integer] The current overall checksum
+  def self.get_current_checksum
+    files = collect_files_to_check
+    cache_file = get_cache_path
+    cache_data = load_combined_cache(cache_file)
+
+    files_lookup = {}
+    cache_data['files'].each { |entry| files_lookup[entry['path']] = entry }
+
+    file_crc32s_with_metadata = calculate_file_checksums(files, files_lookup)
+
+    file_crc32s = file_crc32s_with_metadata.map { |_, meta| meta['crc32'] }.sort
+
+    overall_checksum = calculate_overall_checksum(file_crc32s)
+
+    overall_checksum
+  end
+
+  # Compare the current checksum with the cached checksum
+  # @param [String] current_checksum The calculated checksum for the current state
+  # @param [String] cached_checksum The checksum retrieved from cache
+  # @return [Boolean] True if checksums match, false otherwise
+  def self.checksums_match?(current_checksum, cached_checksum)
+    current_checksum == cached_checksum
+  end
+
+  # Calculate the overall checksum from individual file checksums
+  # @param [Array<Integer>] file_crc32s Array of individual file CRC32 values
+  # @return [Integer] The overall CRC32 as an integer
+  def self.calculate_overall_checksum(file_crc32s)
+    Zlib.crc32(file_crc32s.join(','), 0)
+  end
+
+  # Collect all files that need to be checked for checksums
+  # @return [Array<String>] List of file paths
+  def self.collect_files_to_check
+    # Define the directories to scan for files
+    modules_dir = File.join(Msf::Config.install_root, 'modules', '**', '*')
+    local_modules_dir = File.join(Msf::Config.user_module_directory, '**', '*')
+    lib_dir = File.join(Msf::Config.install_root, 'lib', '**', '*')
+    # Gather all files from the specified directories
+    Dir.glob([modules_dir, lib_dir, local_modules_dir]).select { |f| File.file?(f) }.sort
+  end
+
+  # Calculate checksums for all files, using the cache when possible
+  # @param [Array<String>] files List of file paths to check
+  # @param [Hash] cache Current cache data
+  # @return [Array<Array>] Array of [file_path, metadata] pairs
+  def self.calculate_file_checksums(files, cache)
+    Parallel.map(files, in_threads: Etc.nprocessors * 2) do |file|
+      # Get file metadata (size and last modified time)
+      file_metadata = File.stat(file)
+      cache_entry = cache[file]
+      # Use cached CRC32 if mtime and size match, otherwise recalculate
+      if cache_entry && cache_entry['mtime'] == file_metadata.mtime.to_i && cache_entry['size'] == file_metadata.size
+        crc32 = cache_entry['crc32']
+      else
+        crc32 = File.open(file, 'rb') { |fd| Zlib.crc32(fd.read, 0) }
+      end
+      # Return file and its metadata for later aggregation
+      [file, {
+        'crc32' => crc32,
+        'mtime' => file_metadata.mtime.to_i,
+        'size' => file_metadata.size
+      }]
+    end
+  end
+
+  # Get the path to the cache file
+  # @return [String] Path to the cache file
+  def self.get_cache_path
+    File.join(Msf::Config.config_directory, "store", CacheMetaDataFile)
+  end
+
+  # Load the combined cache from disk (contains both files and checksum)
+  # @param [String] cache_file Path to the cache file
+  # @return [Hash] The loaded cache with 'files' and 'checksum' keys, or empty structure if file doesn't exist
+  def self.load_combined_cache(cache_file)
+    if File.exist?(cache_file)
+      cache_content = JSON.parse(File.read(cache_file))
+      # Ensure the cache has the expected structure
+      {
+        'files' => cache_content['files'] || [],
+        'checksum' => cache_content['checksum']
+      }
+    else
+      { 'files' => [], 'checksum' => nil }
+    end
+  end
+
+  # Save the combined cache to disk (files and checksum in one file)
+  # @param [String] cache_file Path to the cache file
+  # @param [Hash] files_cache The per-file cache data
+  # @param [Integer] overall_checksum The overall checksum
+  # @return [void]
+  def self.save_combined_cache(cache_file, files_cache, overall_checksum)
+    # Ensure the directory for the cache file exists before writing
+    FileUtils.mkdir_p(File.dirname(cache_file))
+
+    cache_content = {
+      'checksum' => overall_checksum,
+      'files' => files_cache
+    }
+
+    File.write(cache_file, JSON.pretty_generate(cache_content))
+  end
+
+  # Get the cached checksum value from the combined cache file
+  # @return [Integer, nil] The cached checksum value or nil if no cache exists
+  def self.get_cached_checksum
+    cache_path = get_cache_path
+    cache_data = load_combined_cache(cache_path)
+    cache_data['checksum']
+  end
+
+  # Update the cache with the current checksum and file data
+  # @param [Integer] current_checksum The current checksum to store in the cache
+  # @return [void]
+  def self.update_cache_checksum(current_checksum)
+    # Recalculate file checksums and update both overall checksum and file cache
+    files = collect_files_to_check
+    cache_file = get_cache_path
+    cache_data = load_combined_cache(cache_file)
+
+    files_lookup = {}
+    cache_data['files'].each { |entry| files_lookup[entry['path']] = entry }
+
+    file_crc32s_with_metadata = calculate_file_checksums(files, files_lookup)
+
+    updated_files_cache = file_crc32s_with_metadata.map do |file_path, metadata|
+      metadata.merge('path' => file_path)
+    end
+
+    updated_files_cache.sort_by! { |entry| entry['path'] }
+
+    # Save both the updated file cache and the new overall checksum
+    save_combined_cache(cache_file, updated_files_cache, current_checksum)
+  end
 end

lib/msf/ui/console/driver.rb

@@ -72,6 +72,16 @@ def initialize(prompt = DefaultPrompt, prompt_char = DefaultPromptChar, opts = {
       elog(e)
     end
 
+    # Check if files have been modified and force immediate loading if so
+    has_modified_metasploit_files = !Msf::Modules::Metadata::Store.valid_checksum?
+
+    if has_modified_metasploit_files
+      current_checksum = Msf::Modules::Metadata::Store.get_current_checksum
+      Msf::Modules::Metadata::Store.update_cache_checksum(current_checksum)
+      # Force immediate module loading when files have changed
+      opts['DeferModuleLoads'] = false
+    end
+
     if opts['DeferModuleLoads'].nil?
       opts['DeferModuleLoads'] = Msf::FeatureManager.instance.enabled?(Msf::FeatureManager::DEFER_MODULE_LOADS)
     end
@@ -163,7 +173,8 @@ def initialize(prompt = DefaultPrompt, prompt_char = DefaultPromptChar, opts = {
       self.framework.init_module_paths(module_paths: opts['ModulePath'], defer_module_loads: opts['DeferModuleLoads'])
     end
 
-    unless opts['DeferModuleLoads']
+    # Refresh module cache if modules are modified, or we're not deferring loads
+    if has_modified_metasploit_files || !opts['DeferModuleLoads']
       framework.threads.spawn("ModuleCacheRebuild", true) do
         framework.modules.refresh_cache_from_module_files
       end