From ea9c28180457e0213c7f63198edd5584c8039907 Mon Sep 17 00:00:00 2001
From: Ben Passmore <contact@passbe.com>
Date: Sat, 30 Nov 2019 15:17:11 +1000
Subject: [PATCH] Updated dns for AR

---
 extensions/dns/dns.rb   | 38 +++++++++++++++++++++++++-------------
 extensions/dns/model.rb | 24 +++++++++++++-----------
 2 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index 8dfe5db3e..6bbbebdfa 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -48,9 +48,10 @@ module BeEF
             pattern = Regexp.new(rule[:pattern], Regexp::IGNORECASE)
             $VERBOSE = verbose
 
-            @database.first_or_create(
-              { :resource => rule[:resource], :pattern => pattern.source },
-              { :response => rule[:response] }
+            @database.find_or_create_by(
+              :resource => rule[:resource].to_s,
+              :pattern => pattern.source,
+              :response => rule[:response]
             ).id
           end
         end
@@ -62,9 +63,11 @@ module BeEF
         # @return [Hash] hash representation of rule (empty hash if rule wasn't found)
         def get_rule(id)
           @lock.synchronize do
-            if is_valid_id?(id)
-              rule = @database.get(id)
-              rule.nil? ? {} : to_hash(rule)
+            begin
+              rule = @database.find(id)
+              return to_hash(rule)
+            rescue ActiveRecord::RecordNotFound
+              return nil
             end
           end
         end
@@ -76,10 +79,15 @@ module BeEF
         # @return [Boolean] true if rule was removed, otherwise false
         def remove_rule!(id)
           @lock.synchronize do
-            if is_valid_id?(id)
-              rule = @database.get(id)
-              rule.nil? ? false : rule.destroy
+            begin
+              rule = @database.find(id)
+              if not rule.nil? and rule.destroy
+                return true
+              end
+            rescue ActiveRecord::RecordNotFound
+              return nil
             end
+            return false
           end
         end
 
@@ -94,14 +102,18 @@ module BeEF
         #
         # @return [Array<Hash>] DNS ruleset (empty array if no rules are currently defined)
         def get_ruleset
-          @lock.synchronize { @database.collect { |rule| to_hash(rule) } }
+          @lock.synchronize { @database.all { |rule| to_hash(rule) } }
         end
 
         # Removes the entire DNS ruleset.
         #
         # @return [Boolean] true if ruleset was destroyed, otherwise false
         def remove_ruleset!
-          @lock.synchronize { @database.destroy }
+          @lock.synchronize do 
+            if @database.destroy_all
+              return true
+            end
+          end
         end
 
         # Starts the DNS server.
@@ -161,7 +173,7 @@ module BeEF
 
             catch (:done) do
               # Find rules matching the requested resource class
-              resources = @database.all(:resource => resource)
+              resources = @database.where(:resource => resource)
               throw :done if resources.length == 0
 
               # Narrow down search by finding a matching pattern
@@ -257,7 +269,7 @@ module BeEF
         #
         # @return [String] resource name stripped of any module/class names
         def format_resource(resource)
-          /::(\w+)$/.match(resource.name)[1]
+          /::(\w+)$/.match(resource)[1]
         end
 
       end
diff --git a/extensions/dns/model.rb b/extensions/dns/model.rb
index befc2ff37..85afab191 100644
--- a/extensions/dns/model.rb
+++ b/extensions/dns/model.rb
@@ -9,26 +9,28 @@ module BeEF
       module Dns
 
         # Represents an individual DNS rule.
-        class Rule < ActiveRecord::Base
-        attribute :id, :String, :key => true
-        attribute :pattern, :Object, :required => true
-        attribute :resource, :Object, :required => true
-        attribute :response, :Object, :required => true
-        attribute :callback, :Object, :required => true
+        class Rule < BeEF::Core::Model
+
           # Hooks the model's "save" event. Validates pattern/response and generates a rule identifier.
-          before_save do |rule|
+          before_save :check_rule
+          self.table_name = 'dns_rule'
+          serialize :response, Array
+
+        private
+
+          def check_rule
             begin
-              validate_pattern(rule.pattern)
-              rule.callback = format_callback(rule.resource, rule.response)
+              validate_pattern(self.pattern)
+              self.callback = format_callback(self.resource.constantize, self.response)
             rescue InvalidDnsPatternError, UnknownDnsResourceError, InvalidDnsResponseError => e
               print_error e.message
               throw :halt
             end
 
-            rule.id = BeEF::Core::Crypto.dns_rule_id
+            #self.id = BeEF::Core::Crypto.dns_rule_id
+
           end
 
-        private
           # Verifies that the given pattern is valid (i.e. non-empty, no null's or printable characters).
           def validate_pattern(pattern)
             raise InvalidDnsPatternError unless BeEF::Filters.is_non_empty_string?(pattern) &&