From 93e0ca7cd5cebf50157c9f4271beec6466bd0bde Mon Sep 17 00:00:00 2001 From: adfoster-r7 Date: Wed, 9 Oct 2024 13:17:48 +0100 Subject: [PATCH 1/4] Improve database module cache performance --- lib/msf/core/db_manager/module_cache.rb | 78 ++++++++++++++++--- .../examples/msf/db_manager/module_cache.rb | 3 +- .../update_all_module_details_refresh.rb | 19 ++--- 3 files changed, 80 insertions(+), 20 deletions(-) diff --git a/lib/msf/core/db_manager/module_cache.rb b/lib/msf/core/db_manager/module_cache.rb index 1c26e3025144..c42d15128016 100644 --- a/lib/msf/core/db_manager/module_cache.rb +++ b/lib/msf/core/db_manager/module_cache.rb @@ -28,7 +28,7 @@ def match_values(values) values.collect { |value| "%#{value}%" } end - def module_to_details_hash(m) + def module_to_details_hash(m, with_mixins: true) res = {} bits = [] @@ -92,8 +92,10 @@ def module_to_details_hash(m) res[:stance] = m.stance.to_s.index("aggressive") ? "aggressive" : "passive" - m.class.mixins.each do |x| - bits << [ :mixin, { :name => x.to_s } ] + if with_mixins + m.class.mixins.each do |x| + bits << [ :mixin, { :name => x.to_s } ] + end end end @@ -269,7 +271,6 @@ def update_all_module_details } Mdm::Module::Detail.find_each do |md| - unless md.ready refresh << md next @@ -291,6 +292,7 @@ def update_all_module_details refresh.each { |md| md.destroy } + new_modules = [] [ ['exploit', framework.exploits], ['auxiliary', framework.auxiliary], @@ -305,14 +307,12 @@ def update_all_module_details next if skip_reference_name_set.include? mn obj = mt[1].create(mn) next if not obj - begin - update_module_details(obj) - rescue ::Exception => e - elog("Error updating module details for #{obj.fullname}", error: e) - end + new_modules <<= obj end end + insert_all(new_modules) + self.framework.cache_initialized = true end @@ -332,7 +332,7 @@ def update_module_details(module_instance) return if not self.migrated ApplicationRecord.connection_pool.with_connection do - info = module_to_details_hash(module_instance) + info = module_to_details_hash(module_instance, with_mixins: false) bits = info.delete(:bits) || [] module_detail = Mdm::Module::Detail.create!(info) @@ -359,4 +359,62 @@ def update_module_details(module_instance) module_detail.save! end end + + private + + # Insert the Msf::Module array into the Mdm::Module::Detail database class + # + # @param [Array] modules + def insert_all(modules) + module_hashes = modules.filter_map do |mod| + begin + hash = module_to_details_hash(mod, with_mixins: false) + # The insert_all API requires all hashes to have the same keys present, so explicitly set these potentially missing keys + hash[:disclosure_date] ||= nil + hash[:default_target] ||= nil + hash[:default_action] ||= nil + hash[:stance] ||= nil + hash + rescue ::Exception => e + elog("Error updating module details for #{mod.fullname}", error: e) + nil + end + end + return if module_hashes.empty? + + # 1) Bulk insert the module detail entries + module_details = module_hashes.map { |mod_hash| mod_hash.except(:bits) } + module_detail_ids = Mdm::Module::Detail.insert_all!(module_details, returning: %w[id]).map { |returning| returning['id'] } + + # 2) Build the hashes for the associations + associations = module_hashes.zip(module_detail_ids).each_with_object(Hash.new { |hash, key| hash[key] = [] }) do |(module_hash, detail_id), acc| + module_hash[:bits].each do |args| + otype, vals = args + + case otype + when :action + acc[Mdm::Module::Action] << { detail_id: detail_id, name: vals[:name] } + when :arch + acc[Mdm::Module::Arch] << { detail_id: detail_id, name: vals[:name] } + when :author + acc[Mdm::Module::Author] << { detail_id: detail_id, name: vals[:name], email: vals[:email] } + when :platform + acc[Mdm::Module::Platform] << { detail_id: detail_id, name: vals[:name] } + when :ref + acc[Mdm::Module::Ref] << { detail_id: detail_id, name: vals[:name] } + when :target + acc[Mdm::Module::Target] << { detail_id: detail_id, index: vals[:index], name: vals[:name] } + end + end + end + + # 3) Insert all of the associations + associations.each do |association_clazz, entries| + next if entries.empty? + + association_clazz.insert_all!(entries) + end + + nil + end end diff --git a/spec/support/shared/examples/msf/db_manager/module_cache.rb b/spec/support/shared/examples/msf/db_manager/module_cache.rb index be886a56d5dc..7cc1f4bcb82e 100644 --- a/spec/support/shared/examples/msf/db_manager/module_cache.rb +++ b/spec/support/shared/examples/msf/db_manager/module_cache.rb @@ -812,7 +812,8 @@ def loader.load_error(module_path, error) allow(db_manager).to receive( :module_to_details_hash ).with( - module_instance + module_instance, + with_mixins: false ).and_return( module_to_details_hash ) diff --git a/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb b/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb index 88a99cb63125..17dbae67d08a 100644 --- a/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb +++ b/spec/support/shared/examples/msf/db_manager/update_all_module_details_refresh.rb @@ -29,19 +29,20 @@ update_all_module_details end - it 'should call update_module_details to create a new Mdm::Module::Detail from the module instance returned by create' do - expect(db_manager).to receive(:update_module_details) do |module_instance| - expect(module_instance).to be_a Msf::Module - expect(module_instance.type).to eq module_detail.mtype - expect(module_instance.refname).to eq module_detail.refname - end - + it 'should create a new Mdm::Module::Detail entry' do update_all_module_details + + aggregate_failures do + expect(Mdm::Module::Detail.count).to eq 1 + db_module_detail = Mdm::Module::Detail.first + expect(db_module_detail.mtype).to eq(module_detail.mtype) + expect(db_module_detail.refname).to eq(module_detail.refname) + end end - context 'with exception raised by #update_module_details' do + context 'with exception raised by #insert_all' do before(:example) do - expect(db_manager).to receive(:update_module_details).and_raise(Exception) + expect(db_manager).to receive(:module_to_details_hash).and_raise(Exception) end it 'should log error' do From dd92e545125c44f7b88cf5e689968467d487801a Mon Sep 17 00:00:00 2001 From: Diana Payton <52059945+oddlittlebird@users.noreply.github.com> Date: Wed, 11 Dec 2024 07:10:04 -0800 Subject: [PATCH 2/4] Update README.md Minor edits to improve the README, added some important information from schema.rb comments. --- db/README.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/db/README.md b/db/README.md index d164fb9c4637..f739837fec80 100644 --- a/db/README.md +++ b/db/README.md @@ -1,2 +1,7 @@ -Contains `modules_metadata_base.json` which contains information about all modules within Metasploit, as well as -`schema.rb` which describes current state of the database schema maintained by Rails ActiveRecord. \ No newline at end of file +This directory contains the following files: + +- `modules_metadata_base.json`, which information about all modules within Metasploit. +- `schema.rb`, which is auto-generated from the current state of the database schema maintained by Rails ActiveRecord. + This file is auto-generated from the current state of the database. + +`schema.rb` is the source Rails uses to define your schema when running `bin/rails db:schema:load`. When creating a new database, `bin/rails db:schema:load` tends to be faster and is potentially less error prone than running all of your migrations from scratch. Old migrations may fail to apply correctly if those migrations use external dependencies or application code. We _strongly_ recommend that you check this file into your version control system. From 25dcd93d0a0e6021c04f3c19ddf05ec8d0d90b8f Mon Sep 17 00:00:00 2001 From: Diana Payton <52059945+oddlittlebird@users.noreply.github.com> Date: Thu, 12 Dec 2024 10:53:57 -0800 Subject: [PATCH 3/4] Update db/README.md Co-authored-by: Arne De Herdt --- db/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/db/README.md b/db/README.md index f739837fec80..2c9a063f6b0b 100644 --- a/db/README.md +++ b/db/README.md @@ -1,6 +1,6 @@ This directory contains the following files: -- `modules_metadata_base.json`, which information about all modules within Metasploit. +- `modules_metadata_base.json`, which contains information about all modules within Metasploit. - `schema.rb`, which is auto-generated from the current state of the database schema maintained by Rails ActiveRecord. This file is auto-generated from the current state of the database. From 051a46a781595da331bc4f94ab46face11064aa8 Mon Sep 17 00:00:00 2001 From: cgranleese-r7 <69522014+cgranleese-r7@users.noreply.github.com> Date: Fri, 13 Dec 2024 16:10:01 +0000 Subject: [PATCH 4/4] Implements feedback --- db/README.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/db/README.md b/db/README.md index 2c9a063f6b0b..82a0198b6689 100644 --- a/db/README.md +++ b/db/README.md @@ -4,4 +4,7 @@ This directory contains the following files: - `schema.rb`, which is auto-generated from the current state of the database schema maintained by Rails ActiveRecord. This file is auto-generated from the current state of the database. -`schema.rb` is the source Rails uses to define your schema when running `bin/rails db:schema:load`. When creating a new database, `bin/rails db:schema:load` tends to be faster and is potentially less error prone than running all of your migrations from scratch. Old migrations may fail to apply correctly if those migrations use external dependencies or application code. We _strongly_ recommend that you check this file into your version control system. +`schema.rb` is the source Rails uses to define your schema when running `bin/rails db:schema:load`. When creating a new +database, `bin/rails db:schema:load` tends to be faster and is potentially less error-prone than running all of your +migrations from scratch. Old migrations may fail to apply correctly if those migrations use external dependencies or +application code. We _strongly_ recommend that you check this file into your version control system.