Land rapid7#4852, vuln note import/export addition

Author: shuckins-r7

lib/msf/core/db_export.rb

@@ -374,6 +374,18 @@ def extract_host_info(report_file)
           report_file.write("      #{el}\n")
         end
 
+        # Notes attached to vulns instead of the host
+        report_file.write("        <notes>\n")
+        @notes.where(vuln_id: e.id).each do |note|
+          report_file.write("      <note>\n")
+          note.attributes.each_pair do |k,v|
+            el = create_xml_element(k,v)
+            report_file.write("      #{el}\n")
+          end
+          report_file.write("      </note>\n")
+        end
+        report_file.write("        </notes>\n")
+
         # References
         report_file.write("        <refs>\n")
         e.refs.each do |ref|

lib/msf/core/db_manager/import/metasploit_framework/xml.rb

@@ -64,6 +64,30 @@ def import_msf_file(args={})
     import_msf_xml(args.merge(:data => data))
   end
 
+  # Imports `Mdm::Note` objects from the XML element.
+  #
+  # @param note [REXML::Element] The Note element
+  # @param allow_yaml [Boolean] whether to allow yaml
+  # @param note_data [Hash] hash containing note attributes to be passed along
+  # @return [void]
+  def import_msf_note_element(note, allow_yaml, note_data={})
+    note_data[:type] = nils_for_nulls(note.elements["ntype"].text.to_s.strip)
+    note_data[:data] = nils_for_nulls(unserialize_object(note.elements["data"], allow_yaml))
+
+    if note.elements["critical"].text
+      note_data[:critical] = true unless note.elements["critical"].text.to_s.strip == "NULL"
+    end
+    if note.elements["seen"].text
+      note_data[:seen] = true unless note.elements["critical"].text.to_s.strip == "NULL"
+    end
+    %W{created-at updated-at}.each { |datum|
+      if note.elements[datum].text
+        note_data[datum.gsub("-","_")] = nils_for_nulls(note.elements[datum].text.to_s.strip)
+      end
+    }
+    report_note(note_data)
+  end
+
   # Imports web_form element using {Msf::DBManager#report_web_form}.
   #
   # @param element [REXML::Element] web_form element.
@@ -280,21 +304,7 @@ def import_msf_xml(args={}, &block)
         note_data = {}
         note_data[:workspace] = wspace
         note_data[:host] = hobj
-        note_data[:type] = nils_for_nulls(note.elements["ntype"].text.to_s.strip)
-        note_data[:data] = nils_for_nulls(unserialize_object(note.elements["data"], allow_yaml))
-
-        if note.elements["critical"].text
-          note_data[:critical] = true unless note.elements["critical"].text.to_s.strip == "NULL"
-        end
-        if note.elements["seen"].text
-          note_data[:seen] = true unless note.elements["critical"].text.to_s.strip == "NULL"
-        end
-        %W{created-at updated-at}.each { |datum|
-          if note.elements[datum].text
-            note_data[datum.gsub("-","_")] = nils_for_nulls(note.elements[datum].text.to_s.strip)
-          end
-        }
-        report_note(note_data)
+        import_msf_note_element(note,allow_yaml,note_data)
       end
 
       host.elements.each('tags/tag') do |tag|
@@ -335,6 +345,13 @@ def import_msf_xml(args={}, &block)
 
         vobj = report_vuln(vuln_data)
 
+        vuln.elements.each("notes/note") do |note|
+          note_data = {}
+          note_data[:workspace] = wspace
+          note_data[:vuln_id] = vobj.id
+          import_msf_note_element(note,allow_yaml,note_data)
+        end
+
         vuln.elements.each("vuln_details/vuln_detail") do |vdet|
           vdet_data = {}
           vdet.elements.each do |det|

lib/msf/core/db_manager/note.rb

@@ -124,6 +124,7 @@ def report_note(opts)
     conditions = { :ntype => ntype }
     conditions[:host_id] = host[:id] if host
     conditions[:service_id] = service[:id] if service
+    conditions[:vuln_id] = opts[:vuln_id]
 
     case mode
     when :unique
@@ -162,6 +163,9 @@ def report_note(opts)
       note.ntype    = ntype
       note.data     = data
     end
+    if opts[:vuln_id]
+      note.vuln_id = opts[:vuln_id]
+    end
     msf_import_timestamps(opts,note)
     note.save!
     ret[:note] = note

lib/msf/core/db_manager/vuln.rb

@@ -45,29 +45,8 @@ def find_vuln_by_details(details_map, host, service=nil)
   end
 
   def find_vuln_by_refs(refs, host, service=nil)
-
-    vuln = nil
-
-    # Try to find an existing vulnerability with the same service & references
-    # If there are multiple matches, choose the one with the most matches
-    if service
-      refs_ids = refs.map{|x| x.id }
-      vuln = service.vulns.where({ 'refs.id' => refs_ids }).includes(:refs).sort { |a,b|
-        ( refs_ids - a.refs.map{|x| x.id } ).length <=> ( refs_ids - b.refs.map{|x| x.id } ).length
-      }.first
-    end
-
-    # Return if we matched based on service
-    return vuln if vuln
-
-    # Try to find an existing vulnerability with the same host & references
-    # If there are multiple matches, choose the one with the most matches
-    refs_ids = refs.map{|x| x.id }
-    vuln = host.vulns.where({ 'service_id' => nil, 'refs.id' => refs_ids }).includes(:refs).sort { |a,b|
-      ( refs_ids - a.refs.map{|x| x.id } ).length <=> ( refs_ids - b.refs.map{|x| x.id } ).length
-    }.first
-
-    return vuln
+    ref_ids = refs.find_all { |ref| ref.name.starts_with? 'CVE-'}
+    host.vulns.joins(:refs).where(service_id: service.try(:id), refs: { id: ref_ids}).first
   end
 
   def get_vuln(wspace, host, service, name, data='')