From 95d0ddbe8725c325bcb81d5d2255bc7e3c1ba454 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Wed, 17 Jul 2013 18:16:46 -0400
Subject: [PATCH] Added new method #remove_ruleset that clears the entire DNS
 ruleset. Included unit tests as well.

---
 extensions/dns/dns.rb          | 19 +++++++++----
 extensions/dns/ruby/rubydns.rb | 52 +++++++++++++++++++---------------
 test/unit/extensions/tc_dns.rb | 10 +++++++
 3 files changed, 53 insertions(+), 28 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index 2318b16e5..9a97bba98 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -83,6 +83,15 @@ module Dns
       @lock.synchronize { @server.remove_rule(id) }
     end
 
+    # Retrieves a specific rule given its id
+    #
+    # @param id [Integer] unique identifier for rule
+    #
+    # @return [Hash] hash representation of rule
+    def get_rule(id)
+      @lock.synchronize { @server.get_rule(id) }
+    end
+
     # Returns an AoH representing the entire current DNS ruleset.
     #
     # Each element is a hash with the following keys:
@@ -97,13 +106,13 @@ module Dns
       @lock.synchronize { @server.get_ruleset }
     end
 
-    # Retrieves a specific rule given its id
+    # Clears the entire DNS ruleset.
     #
-    # @param id [Integer] unique identifier for rule
+    # Requests made after doing so will be passed through to the root nameservers.
     #
-    # @return [Hash] hash representation of rule
-    def get_rule(id)
-      @lock.synchronize { @server.get_rule(id) }
+    # @return [Boolean] true on success, false on failure
+    def remove_ruleset
+      @lock.synchronize { @server.remove_ruleset }
     end
 
     private
diff --git a/extensions/dns/ruby/rubydns.rb b/extensions/dns/ruby/rubydns.rb
index 4e2c795bd..5f54f8fcb 100644
--- a/extensions/dns/ruby/rubydns.rb
+++ b/extensions/dns/ruby/rubydns.rb
@@ -55,6 +55,17 @@ module RubyDNS
 
     end
 
+    # New method that loads all rules from the database at server startup
+    def load_rules
+      BeEF::Core::Models::Dns::Rule.each do |rule|
+        id = rule.id
+        pattern = [rule.pattern, rule.type]
+        block = eval rule.block
+
+        @rules << Rule.new(id, pattern, block)
+      end
+    end
+
     # Now includes BeEF database support and checks for already present rules
     def match(*pattern, block)
       id = ''
@@ -117,15 +128,22 @@ module RubyDNS
       rule != nil ? rule.destroy : false
     end
 
-    # New method that loads all rules from the database at server startup
-    def load_rules
-      BeEF::Core::Models::Dns::Rule.each do |rule|
-        id = rule.id
-        pattern = [rule.pattern, rule.type]
-        block = eval rule.block
+    # New method that returns a hash representing the given rule
+    def get_rule(id)
+      result = {}
 
-        @rules << Rule.new(id, pattern, block)
+      begin
+        rule = BeEF::Core::Models::Dns::Rule.get!(id)
+
+        result[:id] = rule.id
+        result[:pattern] = rule.pattern
+        result[:type] = rule.type.to_s.split('::')[-1]
+        result[:response] = parse_response(rule.block)
+      rescue DataMapper::ObjectNotFoundError => e
+        @logger.error(e.message)
       end
+
+      result
     end
 
     # New method that returns the entire DNS ruleset as an AoH
@@ -146,22 +164,10 @@ module RubyDNS
       result
     end
 
-    # New method that returns a hash representing the given rule
-    def get_rule(id)
-      result = {}
-
-      begin
-        rule = BeEF::Core::Models::Dns::Rule.get!(id)
-
-        result[:id] = rule.id
-        result[:pattern] = rule.pattern
-        result[:type] = rule.type.to_s.split('::')[-1]
-        result[:response] = parse_response(rule.block)
-      rescue DataMapper::ObjectNotFoundError => e
-        @logger.error(e.message)
-      end
-
-      result
+    # New method that removes the entire DNS ruleset
+    def remove_ruleset
+      @rules = []
+      BeEF::Core::Models::Dns::Rule.destroy
     end
 
     private
diff --git a/test/unit/extensions/tc_dns.rb b/test/unit/extensions/tc_dns.rb
index 092518dd2..638ee6fa1 100644
--- a/test/unit/extensions/tc_dns.rb
+++ b/test/unit/extensions/tc_dns.rb
@@ -155,6 +155,7 @@ class TC_Dns < Test::Unit::TestCase
   # Tests the removal of unknown DNS rules
   def test_11_remove_rule_bad
     removed = @@dns.remove_rule(42)
+
     assert(!removed)
   end
 
@@ -172,6 +173,15 @@ class TC_Dns < Test::Unit::TestCase
     check_rule(ruleset[3], {:pattern => 'j.random.hacker', :type => 'A', :response => '4.2.4.2'})
   end
 
+  # Tests the removal of the entire DNS ruleset
+  def test_13_remove_ruleset
+    removed = @@dns.remove_ruleset
+    ruleset = @@dns.get_ruleset
+
+    assert(removed)
+    assert_equal(0, ruleset.length)
+  end
+
   private
 
   # Compares each key in hash 'rule' with the respective key in hash 'expected'