From cdaebc29318d243eb2a3747c4dee51e13dbdd4ff Mon Sep 17 00:00:00 2001
From: HD Moore <hd_moore@rapid7.com>
Date: Sun, 14 Feb 2010 19:34:22 +0000
Subject: [PATCH] Big change to how report_note() works, it now accepts an
 :update parameter that defines whether the data is unique for the host/ntype
 (:unique), unique for the host/ntype/data (:unique_data), or should be
 inserted no matter what (:insert)

git-svn-id: file:///home/svn/framework3/trunk@8495 4d416f70-5f16-0410-b530-b9f4589650da
---
 lib/msf/core/db.rb | 69 ++++++++++++++++++++++++++++++++++------------
 1 file changed, 52 insertions(+), 17 deletions(-)

diff --git a/lib/msf/core/db.rb b/lib/msf/core/db.rb
index ba0e990a35..59312516b9 100644
--- a/lib/msf/core/db.rb
+++ b/lib/msf/core/db.rb
@@ -171,6 +171,10 @@ class DBManager
 		return addr if addr.kind_of? Host
 		wait = opts.delete(:wait)
 
+		if opts[:host_mac]
+			opts[:mac] = opts.delete(:host_mac)
+		end
+
 		ret = {}
 		task = queue( Proc.new {
 			if opts[:comm] and opts[:comm].length > 0
@@ -416,29 +420,59 @@ class DBManager
 			end
 		end
 
+		# Update Modes can be :unique, :unique_data, :insert
+		mode = opts[:update] || :unique
+
 		ret = {}
 		task = queue(Proc.new {
 			if addr and not host
 				host = get_host(addr)
 			end
 
-			ntype = opts.delete(:type) || opts.delete(:ntype) || return
-			data  = opts[:data] || return
+			ntype  = opts.delete(:type) || opts.delete(:ntype) || return
+			data   = opts[:data] || return
+			method = nil
+			args   = []
+			note   = nil
 
-			method = "find_or_initialize_by_ntype_and_data"
-			args = [ ntype, data.to_yaml ]
-			if host
-				method << "_and_host_id"
-				args.push(host.id)
-			end
-			if opts[:service] and opts[:service].kind_of? Service
-				method << "_and_service_id"
-				args.push(opts[:service].id)
+			case mode
+			when :unique
+				method = "find_or_initialize_by_ntype"
+				args = [ ntype ]
+			when :unique_data
+				method = "find_or_initialize_by_ntype_and_data"
+				args = [ ntype, data.to_yaml ]
 			end
 
-			note = workspace.notes.send(method, *args)
-			if (note.changed?)
+			# Find and update a record by type
+			if(method)
+				if host
+					method << "_and_host_id"
+					args.push(host[:id])
+				end
+				if opts[:service] and opts[:service].kind_of? Service
+					method << "_and_service_id"
+					args.push(opts[:service][:id])
+				end
+
+				note = workspace.notes.send(method, *args)
+				if (note.changed?)
+					note.data    = data
+					note.created = Time.now
+					note.save!
+				end
+			# Insert a brand new note record no matter what
+			else
+				note = workspace.notes.new
 				note.created = Time.now
+				if host
+					note.host_id = host[:id]
+				end
+				if opts[:service] and opts[:service].kind_of? Service
+					note.service_id = opts[:service][:id]
+				end
+				note.ntype = ntype
+				note.data  = data
 				note.save!
 			end
 
@@ -480,10 +514,11 @@ class DBManager
 		proto   = proto.downcase
 
 		note = {
-			:ntype => "auth.#{proto}",
-			:host => host,
-			:service => service,
-			:data => opts
+			:ntype    => "auth.#{proto}",
+			:host     => host,
+			:service  => service,
+			:data     => opts,
+			:update   => :unique_data
 		}
 
 		return report_note(note)