Fix grouping of transitive vulns and add spec

Properly merge the contents of the grouping sets when combining sets of
vulns. All permutations of simple double-transitive sets are now tested.

lib/msf/core/analyze.rb

@@ -5,23 +5,32 @@ class Msf::Analyze
   end
 
   def host(eval_host)
-    suggested_modules = []
-
-    mrefs, _mports, _mservs = Msf::Modules::Metadata::Cache.instance.all_exploit_maps
-
     unless eval_host.vulns
       return {}
     end
 
     # Group related vulns
-    vulns = eval_host.vulns.map do |vuln|
+    vuln_families = group_vulns(eval_host.vulns)
+
+    # finds all modules that have references matching those found on host vulns with service data
+    suggested_modules = suggest_modules_for_vulns(eval_host, vuln_families)
+
+    {results: suggested_modules}
+  end
+
+  private
+
+  def group_vulns(vulns)
+    return [] if vulns.empty?
+
+    vulns = vulns.map do |vuln|
       [vuln, Set.new(vuln.refs.map {|r| r.name.upcase})]
     end
     grouped_vulns = Hash.new
 
     vulns.each_index do |ii|
       vuln, refs = vulns[ii]
-      grouped_vulns[vuln] ||= [Array.new, Set.new]
+      grouped_vulns[vuln] ||= [Set.new, Set.new]
       grouped_vulns[vuln][0] << vuln
       grouped_vulns[vuln][1].merge(refs)
 
@@ -29,17 +38,34 @@ class Msf::Analyze
         # TODO: measure if sorting the refs ahead of time and doing a O(n + m)
         # walk here is faster
         if candidate_refs.intersect? refs
-          grouped_vulns[candidate_match] = grouped_vulns[vuln]
+          if grouped_vulns[candidate_match]
+            # For merging two different transitive sets that overlap, we need
+            # to merge the individual grouping elements and then use those
+            # cells inside both the grouping arrays so that all the
+            # already-visited vulns are rolled into the big new set
+            grouped_vulns[candidate_match][0] = grouped_vulns[candidate_match][0].merge(grouped_vulns[vuln][0])
+            grouped_vulns[candidate_match][1] = grouped_vulns[candidate_match][1].merge(grouped_vulns[vuln][1])
+            grouped_vulns[vuln][0] = grouped_vulns[candidate_match][0]
+            grouped_vulns[vuln][1] = grouped_vulns[candidate_match][1]
+          else
+            # Whoever was initialized first has the canonical set
+            grouped_vulns[candidate_match] = grouped_vulns[vuln]
+          end
         end
       end
     end
 
     vuln_families = grouped_vulns.values
-    vuln_families = vuln_families.uniq! || vuln_families
-    # finds all modules that have references matching those found on host vulns with service data
+    vuln_families.uniq! || vuln_families
+  end
+
+  def suggest_modules_for_vulns(eval_host, vuln_families)
+    mrefs, _mports, _mservs = Msf::Modules::Metadata::Cache.instance.all_exploit_maps
+    suggested_modules = []
+
     evaluated_module_targets = Set.new
     to_evaluate_with_defaults = Array.new
-    vuln_families&.each do |vulns, refs|
+    vuln_families.each do |vulns, refs|
       found_modules = mrefs.values_at(*refs).compact.reduce(:+)
       next unless found_modules
 
@@ -79,6 +105,6 @@ class Msf::Analyze
       evaluated_module_targets << [fnd_mod, port]
     end
 
-    {results: suggested_modules}
+    suggested_modules
   end
 end

spec/lib/msf/core/analyze_spec.rb

@@ -0,0 +1,217 @@
+require 'spec_helper'
+
+RSpec.describe Msf::Analyze do
+  context '#group_vulns' do
+    subject(:msf_analyze) { Msf::Analyze.new(nil) }
+    let(:ref_1) { FactoryBot.create(:mdm_ref) }
+    let(:ref_2) { FactoryBot.create(:mdm_ref) }
+    let(:ref_3) { FactoryBot.create(:mdm_ref) }
+    let(:ref_4) { FactoryBot.create(:mdm_ref) }
+
+    let(:vuln_a) { FactoryBot.create(:mdm_vuln) }
+    let(:vuln_ap1) { FactoryBot.create(:mdm_vuln) }
+    let(:vuln_b) { FactoryBot.create(:mdm_vuln) }
+    let(:vuln_c) { FactoryBot.create(:mdm_vuln) }
+    let(:vuln_cta) { FactoryBot.create(:mdm_vuln) }
+    let(:vuln_d) { FactoryBot.create(:mdm_vuln) }
+    let(:vuln_dtc) { FactoryBot.create(:mdm_vuln) }
+
+    let!(:vuln_a_refnames) do
+      refs = [
+        ref_1
+      ]
+      allow(vuln_a).to receive(:refs).and_return(refs)
+      allow(vuln_ap1).to receive(:refs).and_return(refs)
+
+      refs.map { |r| r.name.upcase }
+    end
+
+    let!(:vuln_b_refnames) do
+      refs = [
+        ref_2
+      ]
+      allow(vuln_b).to receive(:refs).and_return(refs)
+
+      refs.map { |r| r.name.upcase }
+    end
+
+    let!(:vuln_c_refnames) do
+      refs = [
+        ref_3
+      ]
+      allow(vuln_c).to receive(:refs).and_return(refs)
+
+      refs.map { |r| r.name.upcase }
+    end
+
+    let!(:vuln_cta_refnames) do
+      refs = [
+        ref_1,
+        ref_3
+      ]
+      allow(vuln_cta).to receive(:refs).and_return(refs)
+
+      refs.map { |r| r.name.upcase }
+    end
+
+    let!(:vuln_d_refnames) do
+      refs = [
+        ref_4
+      ]
+      allow(vuln_d).to receive(:refs).and_return(refs)
+
+      refs.map { |r| r.name.upcase }
+    end
+
+    let!(:vuln_dtc_refnames) do
+      refs = [
+        ref_4,
+        ref_3
+      ]
+      allow(vuln_dtc).to receive(:refs).and_return(refs)
+
+      refs.map { |r| r.name.upcase }
+    end
+
+    it 'should return an Array' do
+      ret = subject.send(:group_vulns, [])
+      expect(ret).to be_an(Array)
+    end
+
+    context 'with one vuln' do
+      subject(:group_vulns) { msf_analyze.send(:group_vulns, [vuln_a]) }
+
+      it 'should return two Sets per vuln family' do
+        expect(subject.size).to be(1)
+        expect(subject.first[0]).to be_a(Set)
+        expect(subject.first[1]).to be_a(Set)
+      end
+
+      it 'should return the vuln' do
+        expect(subject.first[0]).to eql(Set.new([vuln_a]))
+      end
+
+      it 'should return the upcased names of the refs in a set' do
+        expect(subject.first[1]).to eql(Set.new(vuln_a_refnames))
+      end
+    end
+
+    context 'with disjoint vulns' do
+      subject(:group_vulns) { msf_analyze.send(:group_vulns, [vuln_a, vuln_b]) }
+
+      it 'should return two Sets per vuln family' do
+        expect(subject.size).to be(2)
+        subject.each do |family|
+          expect(family[0]).to be_a(Set)
+          expect(family[1]).to be_a(Set)
+        end
+      end
+
+      it 'should return the vulns separately' do
+        expect(subject[0][0]).to eql(Set.new([vuln_a]))
+        expect(subject[1][0]).to eql(Set.new([vuln_b]))
+      end
+
+      it 'should return the upcased names of the refs in separate sets' do
+        expect(subject[0][1]).to eql(Set.new(vuln_a_refnames))
+        expect(subject[1][1]).to eql(Set.new(vuln_b_refnames))
+      end
+    end
+
+    context 'with overlapping vulns' do
+      subject(:group_vulns) { msf_analyze.send(:group_vulns, [vuln_a, vuln_ap1]) }
+
+      it 'should return two Sets per vuln family' do
+        expect(subject.size).to be(1)
+        subject.each do |family|
+          expect(family[0]).to be_a(Set)
+          expect(family[1]).to be_a(Set)
+        end
+      end
+
+      it 'should return the vulns together' do
+        expect(subject[0][0]).to eql(Set.new([vuln_a, vuln_ap1]))
+      end
+
+      it 'should return the upcased names of the refs in a set' do
+        expect(subject[0][1]).to eql(Set.new(vuln_a_refnames))
+      end
+    end
+
+    context 'with overlapping and disjoint vulns' do
+      subject(:group_vulns) { msf_analyze.send(:group_vulns, [vuln_a, vuln_b, vuln_ap1]) }
+
+      it 'should return two Sets per vuln family' do
+        expect(subject.size).to be(2)
+        subject.each do |family|
+          expect(family[0]).to be_a(Set)
+          expect(family[1]).to be_a(Set)
+        end
+      end
+
+      it 'should return the vulns with the same references together' do
+        expect(subject[0][0]).to eql(Set.new([vuln_a, vuln_ap1]))
+        expect(subject[1][0]).to eql(Set.new([vuln_b]))
+      end
+
+      it 'should return the upcased names of the refs separate sets' do
+        expect(subject[0][1]).to eql(Set.new(vuln_a_refnames))
+        expect(subject[1][1]).to eql(Set.new(vuln_b_refnames))
+      end
+    end
+
+    context 'with transitive vulns' do
+      %w(vuln_a vuln_c vuln_cta).permutation do |perm|
+        context "in permutation #{perm.inspect}" do
+          # One the one hand, we need to test all these permutations, on the
+          # other I'm sorry.
+          let(:vuln_permutation) { eval("[#{perm.join(',')}]") }
+          subject(:group_vulns) { msf_analyze.send(:group_vulns, vuln_permutation) }
+
+          it 'should return two Sets per vuln family' do
+            expect(subject.size).to be(1)
+            subject.each do |family|
+              expect(family[0]).to be_a(Set)
+              expect(family[1]).to be_a(Set)
+            end
+          end
+
+          it 'should return the vulns together' do
+            expect(subject[0][0]).to eql(Set.new(vuln_permutation))
+          end
+
+          it 'should return the upcased names of the refs a set' do
+            expect(subject[0][1]).to eql(Set.new(vuln_a_refnames.concat(vuln_c_refnames)))
+          end
+        end
+      end
+    end
+
+    context 'with double-transitive vulns' do
+      %w(vuln_a vuln_c vuln_cta vuln_d vuln_dtc).permutation do |perm|
+        context "in permutation #{perm.inspect}" do
+          # One the one hand, we need to test all these permutations, on the
+          # other I'm sorry.
+          let(:vuln_permutation) { eval("[#{perm.join(',')}]") }
+          subject(:group_vulns) { msf_analyze.send(:group_vulns, vuln_permutation) }
+
+          it 'should return two Sets per vuln family' do
+            expect(subject.size).to be(1)
+            subject.each do |family|
+              expect(family[0]).to be_a(Set)
+              expect(family[1]).to be_a(Set)
+            end
+          end
+
+          it 'should return the vulns together' do
+            expect(subject[0][0]).to eql(Set.new(vuln_permutation))
+          end
+
+          it 'should return the upcased names of the refs a set' do
+            expect(subject[0][1]).to eql(Set.new(vuln_a_refnames.concat(vuln_c_refnames, vuln_d_refnames)))
+          end
+        end
+      end
+    end
+  end
+end