From d05397e0a9926251ad94a08f2e7b42efa9c8d764 Mon Sep 17 00:00:00 2001
From: Brendan Coles <bcoles@gmail.com>
Date: Sun, 28 Jun 2015 17:22:16 +0000
Subject: [PATCH] Add NetworkService.add and NetworkHost.add

---
 extensions/network/models/network_host.rb    | 37 ++++++++++++++++++++
 extensions/network/models/network_service.rb | 35 ++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/extensions/network/models/network_host.rb b/extensions/network/models/network_host.rb
index 0de5ce0d4..3ac0dd7e6 100644
--- a/extensions/network/models/network_host.rb
+++ b/extensions/network/models/network_host.rb
@@ -24,6 +24,43 @@ module BeEF
         property :mac, String, :lazy => false
         property :cid, String, :lazy => false # command id or 'init'
 
+        #
+        # Stores a network host in the data store
+        #
+        def self.add(host={})
+          (print_error "Invalid hooked browser session"; return) unless BeEF::Filters.is_valid_hook_session_id?(host[:hooked_browser_id])
+          (print_error "Invalid IP address"; return) unless BeEF::Filters.is_valid_ip?(host[:ip])
+
+          # 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[:hostname].nil? && host[:type].nil? && host[:os].nil? && host[:mac].nil?
+            return unless BeEF::Core::Models::NetworkHost.all(
+              :hooked_browser_id => host[:hooked_browser_id],
+              :ip => host[:ip]).empty?
+          end
+
+          # 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],
+            :cid => host[:cid])
+          result = network_host.save
+          (print_error "Failed to save network host"; return) if result.nil?
+
+          network_host
+        end
+
       end
 
     end
diff --git a/extensions/network/models/network_service.rb b/extensions/network/models/network_service.rb
index c806b2aa1..724c4e5f7 100644
--- a/extensions/network/models/network_service.rb
+++ b/extensions/network/models/network_service.rb
@@ -23,6 +23,41 @@ module BeEF
         property :type, String, :lazy => false
         property :cid, String, :lazy => false # command id or 'init'
 
+        #
+        # Stores a network service in the data store
+        #
+        def self.add(service={})
+          (print_error "Invalid hooked browser session"; return) if not BeEF::Filters.is_valid_hook_session_id?(service[:hooked_browser_id])
+          (print_error "Invalid IP address"; return) if not BeEF::Filters.is_valid_ip?(service[:ip])
+
+          # store the returned network host details
+          BeEF::Core::Models::NetworkHost.add(
+            :hooked_browser_id => service[:hooked_browser_id],
+            :ip => service[:ip],
+            :cid => service[:cid])
+
+          # prevent duplicates
+          return unless BeEF::Core::Models::NetworkService.all(
+            :hooked_browser_id => service[:hooked_browser_id],
+            :proto => service[:proto],
+            :ip => service[:ip],
+            :port => service[:port],
+            :type => service[:type]).empty?
+
+          # store the returned network service details
+          network_service = BeEF::Core::Models::NetworkService.new(
+            :hooked_browser_id => service[:hooked_browser_id],
+            :proto => service[:proto],
+            :ip => service[:ip],
+            :port => service[:port],
+            :type => service[:type],
+            :cid => service[:cid])
+          result = network_service.save
+          (print_error "Failed to save network service"; return) if result.nil?
+
+          network_service
+        end
+
       end
 
     end