From 96e4063e2ad1a28009167a447e0f28abc39255a1 Mon Sep 17 00:00:00 2001
From: Ben Passmore <contact@passbe.com>
Date: Sun, 3 Nov 2019 08:02:52 +1000
Subject: [PATCH] AR corrected models

---
 core/main/models/browserdetails.rb | 13 +++++++------
 core/main/models/command.rb        |  8 +++++---
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/core/main/models/browserdetails.rb b/core/main/models/browserdetails.rb
index 907e90031..dc35444c8 100644
--- a/core/main/models/browserdetails.rb
+++ b/core/main/models/browserdetails.rb
@@ -17,7 +17,7 @@ module Models
     # Returns the requested value from the data store
     #
     def self.get(session_id, key) 
-      browserdetail = first(:session_id => session_id, :detail_key => key)
+      browserdetail = self.where(:session_id => session_id, :detail_key => key).first
       
       return nil if browserdetail.nil?
       return nil if browserdetail.detail_value.nil?
@@ -28,19 +28,20 @@ module Models
     # Stores or updates an existing key->value pair in the data store
     #
     def self.set(session_id, detail_key, detail_value) 
-      browserdetails = BeEF::Core::Models::BrowserDetails.all(
+      browserdetails = BeEF::Core::Models::BrowserDetails.where(
         :session_id => session_id,
-        :detail_key => detail_key )
-      if browserdetails.nil? || browserdetails.empty?
+        :detail_key => detail_key ).first
+      if browserdetails.nil?
         # store the new browser details key/value
         browserdetails = BeEF::Core::Models::BrowserDetails.new(
               :session_id   => session_id,
               :detail_key   => detail_key,
               :detail_value => detail_value || '')
-        result = browserdetails.save
+        result = browserdetails.save!
       else
         # update the browser details key/value
-        result = browserdetails.update(:detail_value => detail_value || '')
+        browserdetails.detail_value = detail_value || ''
+        result = browserdetails.save!
         print_debug "Browser has updated '#{detail_key}' to '#{detail_value}'"
       end
 
diff --git a/core/main/models/command.rb b/core/main/models/command.rb
index 995e4ecc4..3a621483a 100644
--- a/core/main/models/command.rb
+++ b/core/main/models/command.rb
@@ -12,6 +12,8 @@ module Models
   class Command < BeEF::Core::Model
 
     has_many :results
+    has_one :command_module
+    has_one :hooked_browser
 
     #
     # Save results and flag that the command has been run on the hooked browser
@@ -30,12 +32,12 @@ module Models
       raise TypeError, '"status" needs to be an integer' unless status.integer?
 
       # @note get the hooked browser structure and id from the database
-      hooked_browser = BeEF::Core::Models::HookedBrowser.first(:session => hook_session_id) || nil
+      hooked_browser = BeEF::Core::Models::HookedBrowser.where(:session => hook_session_id).first || nil
       raise TypeError, "hooked_browser is nil" if hooked_browser.nil?
       raise TypeError, "hooked_browser.id is nil" if hooked_browser.id.nil?
 
       # @note get the command module data structure from the database
-      command = first(:id => command_id, :hooked_browser_id => hooked_browser.id) || nil
+      command = self.where(:id => command_id, :hooked_browser_id => hooked_browser.id).first || nil
       raise TypeError, "command is nil" if command.nil?
 
       # @note create the entry for the results 
@@ -45,7 +47,7 @@ module Models
         :status => status,
         :date => Time.now.to_i
       )
-      command.save
+      command.save!
 
       s = show_status(status)
       log = "Hooked browser [id:#{hooked_browser.id}, ip:#{hooked_browser.ip}]"