From 6ba107469507be5fdb3824e22d81e8b1ead7dde4 Mon Sep 17 00:00:00 2001
From: Brendan Coles <bcoles@gmail.com>
Date: Wed, 23 Dec 2015 11:14:55 +0000
Subject: [PATCH] Prevent duplicates

---
 extensions/network/models/network_host.rb | 36 ++++++++++-------------
 1 file changed, 15 insertions(+), 21 deletions(-)

diff --git a/extensions/network/models/network_host.rb b/extensions/network/models/network_host.rb
index a975faaa6..7c6fca3d7 100644
--- a/extensions/network/models/network_host.rb
+++ b/extensions/network/models/network_host.rb
@@ -39,29 +39,23 @@ module BeEF
             end
           end
 
-          # update lastseen
-          BeEF::Core::Models::NetworkHost.all(
-            :hooked_browser_id => host[:hooked_browser_id],
-            :ip => host[:ip]).update( :lastseen => Time.now )
+          # prepare new host for data store
+          new_host = {}
+          new_host[:hooked_browser_id] = host[:hooked_browser_id]
+          new_host[:ip] = host[:ip]
+          new_host[:hostname] = host[:hostname] unless host[:hostname].nil?
+          new_host[:type] = host[:type] unless host[:type].nil?
+          new_host[:os] = host[:os] unless host[:os].nil?
+          new_host[:mac] = host[:mac] unless host[:mac].nil?
 
-          # prevent duplicates
-          return unless BeEF::Core::Models::NetworkHost.all(
-            :hooked_browser_id => host[:hooked_browser_id],
-            :ip => host[:ip],
-            :hostname => host[:hostname],
-            :type => host[:type],
-            :os => host[:os],
-            :mac => host[:mac]).empty?
+          # if host already exists in data store with the same details
+          # then update lastseen and return
+          existing_host = BeEF::Core::Models::NetworkHost.all(new_host)
+          (existing_host.update( :lastseen => Time.now ); return) unless existing_host.empty?
 
-          # store the returned network host details
-          network_host = BeEF::Core::Models::NetworkHost.new(
-            :hooked_browser_id => host[:hooked_browser_id],
-            :ip => host[:ip],
-            :hostname => host[:hostname],
-            :type => host[:type],
-            :os => host[:os],
-            :mac => host[:mac],
-            :lastseen => Time.now)
+          # store the new network host details
+          new_host[:lastseen] = Time.now
+          network_host = BeEF::Core::Models::NetworkHost.new(new_host)
           result = network_host.save
           (print_error "Failed to save network host"; return) if result.nil?