diff --git a/extensions/network/models/network_host.rb b/extensions/network/models/network_host.rb index ffdcf8476..d8d0886f9 100644 --- a/extensions/network/models/network_host.rb +++ b/extensions/network/models/network_host.rb @@ -9,15 +9,9 @@ module BeEF # # Table stores each host identified on the zombie browser's network(s) # - class NetworkHost < ActiveRecord::Base - attribute :id, :Serial - attribute :hooked_browser_id, :Text, lazy: false - attribute :ip, :Text, lazy: false - attribute :hostname, :String, lazy: false - attribute :type, :String, lazy: false # proxy, router, gateway, dns, etc - attribute :os, :String, lazy: false - attribute :mac, :String, lazy: false - attribute :lastseen, :String, length: 15 + class NetworkHost < BeEF::Core::Model + belongs_to :hooked_browser + # # Stores a network host in the data store # @@ -45,23 +39,24 @@ module BeEF 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[:ntype] = host[:ntype] unless host[:ntype].nil? new_host[:os] = host[:os] unless host[:os].nil? new_host[:mac] = host[:mac] unless host[:mac].nil? # 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 = BeEF::Core::Models::NetworkHost.where(hooked_browser_id: new_host[:hooked_browser_id], ip: new_host[:ip]).limit(1) unless existing_host.empty? - existing_host.update(lastseen: Time.new.to_i) + existing_host = existing_host.first + existing_host.lastseen = Time.new.to_i + existing_host.save! return end # store the new network host details new_host[:lastseen] = Time.new.to_i network_host = BeEF::Core::Models::NetworkHost.new(new_host) - result = network_host.save - if result.nil? + if network_host.save print_error 'Failed to save network host' return end @@ -78,7 +73,7 @@ module BeEF return end - host = BeEF::Core::Models::NetworkHost.get(id.to_i) + host = BeEF::Core::Models::NetworkHost.find(id.to_i) if host.nil? print_error "Failed to remove network host [id: #{id}]. Host does not exist." return @@ -95,7 +90,7 @@ module BeEF hooked_browser_id: hooked_browser_id, ip: ip, hostname: hostname, - type: type, + ntype: ntype, os: os, mac: mac, lastseen: lastseen diff --git a/extensions/network/models/network_service.rb b/extensions/network/models/network_service.rb index fb64b39a7..af8b7dea1 100644 --- a/extensions/network/models/network_service.rb +++ b/extensions/network/models/network_service.rb @@ -9,15 +9,11 @@ module BeEF # # Table stores each open port identified on the zombie browser's network(s) # - class NetworkService < ActiveRecord::Base - - attribute :id, :Serial - attribute :hooked_browser_id, :Text, lazy: false - attribute :proto, :String, lazy: false - attribute :ip, :Text, lazy: false - attribute :port, :String, lazy: false - attribute :type, :String, lazy: false + class NetworkService < BeEF::Core::Model + belongs_to :hooked_browser + + # # Stores a network service in the data store # def self.add(service = {}) @@ -44,19 +40,20 @@ module BeEF end # store the returned network host details - BeEF::Core::Models::NetworkHost.add( + BeEF::Core::Models::NetworkHost.create( hooked_browser_id: service[:hooked_browser_id], ip: service[:ip] ) # prevent duplicates - return unless BeEF::Core::Models::NetworkService.all( + total = BeEF::Core::Models::NetworkService.where( hooked_browser_id: service[:hooked_browser_id], proto: service[:proto], ip: service[:ip], port: service[:port], - type: service[:type] - ).empty? + ntype: service[:ntype] + ).length + return if total > 0 # store the returned network service details network_service = BeEF::Core::Models::NetworkService.new( @@ -64,10 +61,9 @@ module BeEF proto: service[:proto], ip: service[:ip], port: service[:port], - type: service[:type] + ntype: service[:ntype] ) - result = network_service.save - if result.nil? + if network_service.save print_error 'Failed to save network service' return end @@ -83,7 +79,7 @@ module BeEF proto: proto, ip: ip, port: port, - type: type + ntype: ntype } end end diff --git a/extensions/network/rest/network.rb b/extensions/network/rest/network.rb index dbd35e6c5..2586d69aa 100644 --- a/extensions/network/rest/network.rb +++ b/extensions/network/rest/network.rb @@ -27,7 +27,7 @@ module BeEF # Returns the entire list of network hosts for all zombies get '/hosts' do begin - hosts = @nh.all(unique: true, order: [:id.asc]) + hosts = @nh.all.distinct.order(:id) count = hosts.length result = {} @@ -47,7 +47,7 @@ module BeEF # Returns the entire list of network services for all zombies get '/services' do begin - services = @ns.all(unique: true, order: [:id.asc]) + services = @ns.all.distinct.order(:id) count = services.length result = {} @@ -69,7 +69,7 @@ module BeEF begin id = params[:id] - hosts = @nh.all(hooked_browser_id: id, unique: true, order: [:id.asc]) + hosts = @nh.where(hooked_browser_id: id).distinct.order(:id) count = hosts.length result = {} @@ -94,7 +94,7 @@ module BeEF begin id = params[:id] - services = @ns.all(hooked_browser_id: id, unique: true, order: [:id.asc]) + services = @ns.where(hooked_browser_id: id).distinct.order(:id) count = services.length result = {} @@ -119,7 +119,7 @@ module BeEF begin id = params[:id] - host = @nh.all(id: id) + host = @nh.find(id) raise InvalidParamError, 'id' if host.nil? halt 404 if host.empty? @@ -139,7 +139,7 @@ module BeEF id = params[:id] raise InvalidParamError, 'id' unless BeEF::Filters.nums_only?(id) - host = @nh.all(id: id) + host = @nh.find(id) halt 404 if host.nil? result = {} @@ -159,7 +159,7 @@ module BeEF begin id = params[:id] - service = @ns.all(id: id) + service = @ns.find(id) raise InvalidParamError, 'id' if service.nil? halt 404 if service.empty?