Fix bug when updating Mdm::Vuln.refs

lib/msf/core/db_manager/ref.rb

@@ -5,6 +5,7 @@ module Msf::DBManager::Ref
   def find_or_create_ref(opts)
     ret = {}
     ret[:ref] = get_ref(opts[:name])
+    ret[:ref] = Mdm::Ref.find_by_id(opts[:id]) if opts[:id]
     return ret[:ref] if ret[:ref]
 
   ::ActiveRecord::Base.connection_pool.with_connection {
@@ -38,4 +39,4 @@ module Msf::DBManager::Ref
     Mdm::Ref.find_by_name(name)
   }
   end
-end
+end

lib/msf/core/db_manager/vuln.rb

@@ -262,8 +262,9 @@ module Msf::DBManager::Vuln
   ::ActiveRecord::Base.connection_pool.with_connection {
     wspace = Msf::Util::DBManager.process_opts_workspace(opts, framework, false)
     opts[:workspace] = wspace if wspace
-    id = opts.delete(:id)
-    Mdm::Vuln.update(id, opts)
+    v = Mdm::Vuln.find(opts.delete(:id))
+    v.update!(opts)
+    v
   }
   end
 

lib/msf/core/web_services/servlet/vuln_servlet.rb

@@ -1,5 +1,7 @@
 module VulnServlet
 
+  JSON_INCLUDES = [:host, :vulns_refs, :refs, :module_refs]
+
   def self.api_path
     '/api/v1/vulns'
   end
@@ -25,9 +27,8 @@ module VulnServlet
       begin
         sanitized_params = sanitize_params(params, env['rack.request.query_hash'])
         data = get_db.vulns(sanitized_params)
-        includes = [:host, :vulns_refs, :refs, :module_refs]
         data = data.first if is_single_object?(data, sanitized_params)
-        set_json_data_response(response: data, includes: includes)
+        set_json_data_response(response: data, includes: JSON_INCLUDES)
       rescue => e
         print_error_and_create_response(error: e, message: 'There was an error retrieving vulns:', code: 500)
       end
@@ -51,8 +52,16 @@ module VulnServlet
         opts = parse_json_request(request, false)
         tmp_params = sanitize_params(params)
         opts[:id] = tmp_params[:id] if tmp_params[:id]
+        # update_vuln
+        if opts[:refs]
+          refs = []
+          opts[:refs].each do |r|
+            refs << get_db.find_or_create_ref(r)
+          end
+          opts[:refs] = refs
+        end
         data = get_db.update_vuln(opts)
-        set_json_data_response(response: data)
+        set_json_data_response(response: data, includes: JSON_INCLUDES)
       rescue => e
         print_error_and_create_response(error: e, message: 'There was an error updating the vuln:', code: 500)
       end
@@ -65,11 +74,11 @@ module VulnServlet
       begin
         opts = parse_json_request(request, false)
         data = get_db.delete_vuln(opts)
-        set_json_data_response(response: data)
+        set_json_data_response(response: data, includes: JSON_INCLUDES)
       rescue => e
         print_error_and_create_response(error: e, message: 'There was an error deleting the vulns:', code: 500)
       end
     }
   end
 
-end
+end