From 56c686de642ea16e6c6b371e6e0f312c030de5e4 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Wed, 23 Apr 2014 12:42:39 -0400
Subject: [PATCH] Implemented #get_ruleset method.

Also refactored #get_rule to use new #to_hash helper method since
---
 extensions/dns/dns.rb | 43 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index 2cf90e66d..b0a7dceeb 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -57,22 +57,13 @@ module BeEF
 
         # Retrieves a specific rule given its identifier.
         #
-        # @param id [Integer] unique identifier for rule
+        # @param id [String] unique identifier for rule
         #
         # @return [Hash] hash representation of rule (empty hash if rule wasn't found)
         def get_rule(id)
           @lock.synchronize do
             rule = BeEF::Core::Models::Dns::Rule.get(id)
-            hash = {}
-
-            unless rule.nil?
-              hash[:id] = rule.id
-              hash[:pattern] = rule.pattern
-              hash[:resource] = rule.resource
-              hash[:response] = rule.response
-            end
-
-            hash
+            rule.nil? ? {} : to_hash(rule)
           end
         end
 
@@ -88,6 +79,20 @@ module BeEF
           end
         end
 
+        # Returns an AoH representing the entire current DNS ruleset.
+        #
+        # Each element is a hash with the following keys:
+        #
+        # * <code>:id</code>
+        # * <code>:pattern</code>
+        # * <code>:resource</code>
+        # * <code>:response</code>
+        #
+        # @return [Array<Hash>] DNS ruleset (empty array if no rules are currently defined)
+        def get_ruleset
+          @lock.synchronize { BeEF::Core::Models::Dns::Rule.collect { |rule| to_hash(rule) } }
+        end
+
         # Entry point for processing incoming DNS requests. Attempts to find a matching rule and
         # sends back its associated response.
         #
@@ -100,6 +105,22 @@ module BeEF
           end
         end
 
+      private
+        # Helper method that converts a DNS rule to a hash.
+        #
+        # @param rule [BeEF::Core::Models::Dns::Rule] rule to be converted
+        #
+        # @return [Hash] hash representation of DNS rule
+        def to_hash(rule)
+          hash = {}
+          hash[:id] = rule.id
+          hash[:pattern] = rule.pattern
+          hash[:resource] = rule.resource
+          hash[:response] = rule.response
+
+          hash
+        end
+
       end
 
     end