Adds cache invalidation logic

Author: cgranleese-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 = 'cache_metadata_base.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,213 @@ def load_cache_from_file_store
     }
   end
 
+  # This method uses a per-file CRC32 cache to avoid recalculating checksums for files that have not changed.
+  # It loads the cache, checks each file's mtime and size, and only recalculates the CRC32 if needed.
+  #
+  # @return [Boolean] True if the current checksum matches the cached one
+  def self.valid_checksum?
+    current_checksum = get_current_checksum
+
+    get_store_cache_path
+    ensure_cache_file_exists(current_checksum)
+
+    cached_sha = get_cached_checksum
+
+    checksums_match?(current_checksum, cached_sha)
+  end
+
+  # Calculate the current checksum for all module and library files
+  # This calculates checksums for each file, caches them, and then
+  # generates an overall checksum from the individual file checksums.
+  #
+  # @return [String] The current overall checksum
+  def self.get_current_checksum
+    files = collect_files_to_check
+    per_file_cache_file = get_per_file_cache_path
+    per_file_cache = load_per_file_cache(per_file_cache_file)
+
+    file_crc32s_with_metadata = calculate_file_checksums(files, per_file_cache)
+
+    updated_cache = file_crc32s_with_metadata.to_h
+    file_crc32s = file_crc32s_with_metadata.map { |_, meta| meta['crc32'] }
+
+    save_per_file_cache(per_file_cache_file, updated_cache)
+
+    calculate_overall_checksum(file_crc32s)
+  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 [String] The hexadecimal representation of the overall CRC32
+  def self.calculate_overall_checksum(file_crc32s)
+    Zlib.crc32(file_crc32s.join).to_s(16)
+  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) }
+      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 per-file cache
+  # @return [String] Path to the per-file cache
+  def self.get_per_file_cache_path
+    File.join(Msf::Config.config_directory, 'store', 'per_file_metadata_cache.json')
+  end
+
+  # Get the path to the cache store file
+  # @return [String] Path to the cache store file
+  def self.get_store_cache_path
+    File.join(Msf::Config.config_directory, "store", CacheMetaDataFile)
+  end
+
+  # Get the path to the DB cache file
+  # @return [String] Path to the DB cache file
+  def self.get_db_cache_path
+    File.join(Msf::Config.install_root, "db", CacheMetaDataFile)
+  end
+
+  # Load the per-file cache from disk
+  # @param [String] cache_file Path to the cache file
+  # @return [Hash] The loaded cache or an empty hash if the file doesn't exist
+  def self.load_per_file_cache(cache_file)
+    File.exist?(cache_file) ? JSON.parse(File.read(cache_file)) : {}
+  end
+
+  # Save the updated per-file cache to disk
+  # @param [String] cache_file Path to the cache file
+  # @param [Hash] updated_cache The cache data to save
+  # @return [void]
+  def self.save_per_file_cache(cache_file, updated_cache)
+    # Ensure the directory for the cache file exists before writing
+    FileUtils.mkdir_p(File.dirname(cache_file))
+    # Save the updated per-file cache to disk
+    File.write(cache_file, JSON.pretty_generate(updated_cache))
+  end
+
+  # Create or update a cache file with the given checksum
+  # @param [String] file_path Path to the cache file
+  # @param [String] checksum The checksum to store
+  # @return [void]
+  def self.create_or_update_cache_file(file_path, checksum)
+    # Ensure directory exists
+    FileUtils.mkdir_p(File.dirname(file_path))
+
+    if File.exist?(file_path)
+      # Update existing file
+      cache_content = JSON.parse(File.read(file_path))
+      cache_content['checksum']['crc32'] = checksum
+    else
+      # Create new file
+      cache_content = {
+        "checksum" => {
+          "crc32" => checksum
+        }
+      }
+    end
+
+    File.write(file_path, JSON.pretty_generate(cache_content))
+  end
+
+  # Ensure the db cache file exists, creating it if necessary
+  # @param [String] current_checksum The current checksum to use if creating a new cache file
+  # @return [void]
+  def self.ensure_cache_file_exists(current_checksum)
+    # Path to the DB cache file
+    cache_db_path = get_db_cache_path
+
+    # Only create the db cache file if it doesn't exist
+    # The user's cache file (~/.msf4/store/cache_metadata_base.json) should only be created when changes are made
+    unless File.exist?(cache_db_path)
+      # Ensure directory exists
+      FileUtils.mkdir_p(File.dirname(cache_db_path))
+      cache_content = {
+        "checksum" => {
+          "crc32" => current_checksum
+        }
+      }
+      File.write(cache_db_path, JSON.pretty_generate(cache_content))
+    end
+  end
+
+  # Get the cached checksum value without creating any new files
+  # @return [String, nil] The cached checksum value or nil if no cache exists
+  def self.get_cached_checksum
+    cache_store_path = get_store_cache_path
+    cache_db_path = get_db_cache_path
+
+    # First try user's cache file
+    if File.exist?(cache_store_path)
+      cache_content = JSON.parse(File.read(cache_store_path))
+      return cache_content.dig('checksum', 'crc32')
+    end
+
+    # Fall back to db cache file
+    if File.exist?(cache_db_path)
+      cache_content = JSON.parse(File.read(cache_db_path))
+      return cache_content.dig('checksum', 'crc32')
+    end
+
+    # If neither exists, return nil to trigger a cache rebuild
+    # This allows the build process to work with neither file present
+    nil
+  end
+
+  # Update the cache checksum file with the current crc32 checksum of the module paths.
+  #
+  # @param [String] current_checksum The current checksum to store in the cache
+  # @return [void]
+  def self.update_cache_checksum(current_checksum)
+    cache_store_path = get_store_cache_path
+    cache_db_path = get_db_cache_path
+
+    if File.exist?(cache_store_path)
+      # Update the existing user cache file
+      create_or_update_cache_file(cache_store_path, current_checksum)
+    elsif File.exist?(cache_db_path)
+      # Copy the DB cache file to the user's directory and update it
+      FileUtils.cp(cache_db_path, cache_store_path)
+      create_or_update_cache_file(cache_store_path, current_checksum)
+    else
+      # Create a new cache file if neither exists
+      create_or_update_cache_file(cache_store_path, current_checksum)
+    end
+  end
 end

lib/msf/ui/console/driver.rb

@@ -163,6 +163,18 @@ def initialize(prompt = DefaultPrompt, prompt_char = DefaultPromptChar, opts = {
       self.framework.init_module_paths(module_paths: opts['ModulePath'], defer_module_loads: opts['DeferModuleLoads'])
     end
 
+    # If the module cache is invalid, we need to invalidate it and update the cache.
+    # We will also need to set 'DeferModuleLoads' to false so that the modules are reloaded.
+    unless Msf::Modules::Metadata::Store.valid_checksum?
+      # Get the current checksum and update the cache with it
+      current_checksum = Msf::Modules::Metadata::Store.get_current_checksum
+      Msf::Modules::Metadata::Store.update_cache_checksum(current_checksum)
+
+      framework.threads.spawn("ModuleCacheRebuild", true) do
+        framework.modules.refresh_cache_from_module_files
+      end
+    end
+
     unless opts['DeferModuleLoads']
       framework.threads.spawn("ModuleCacheRebuild", true) do
         framework.modules.refresh_cache_from_module_files

tools/automation/cache/build_new_cache.sh

@@ -1,6 +1,7 @@
 #!/bin/sh -ex
 bundle install
 rm -f db/modules_metadata_base.json
+rm -f db/cache_metadata_base.json
 git ls-files modules/ -z | xargs -0 -n1 -P `nproc` -I{} -- git log -1 --format="%ai {}" {} | while read -r udate utime utz ufile ; do
   touch -d "$udate $utime" $ufile
 done