From 054767c898e6f47348c2b9211ea333ba198689e3 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Fri, 17 May 2013 23:02:40 -0400
Subject: [PATCH] Added RESTful API route for /api/dns/rule/:id. This will
 return a single rule given its unique id.

---
 extensions/dns/dns.rb          | 19 +++++++++++---
 extensions/dns/rest/dns.rb     | 46 +++++++++++-----------------------
 extensions/dns/ruby/rubydns.rb | 18 +++++++++++++
 3 files changed, 47 insertions(+), 36 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index 0be4a963c..8580fa02b 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -17,11 +17,11 @@ module DNS
 
     include Singleton
 
-    # @!method instance
+    # @!method self.instance
     #  Returns the singleton instance. Use this in place of {#initialize}.
 
-    # @note This method cannot be invoked! Use {#instance} instead.
-    # @see #instance
+    # @note This method cannot be invoked! Use {.instance} instead.
+    # @see ::instance
     def initialize
       @lock = Mutex.new
       @server = nil
@@ -100,7 +100,7 @@ module DNS
     # * <code>:block</code>
     #
     # @return [Array<Hash>] DNS ruleset (empty if no rules are currently loaded)
-    def get_rules
+    def get_ruleset
       @lock.synchronize do
         result = []
 
@@ -119,6 +119,17 @@ module DNS
       end
     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 do
+        @server.get_rule(id)
+      end
+    end
+
   end
 
 end
diff --git a/extensions/dns/rest/dns.rb b/extensions/dns/rest/dns.rb
index 74fc69c15..dcf7e0849 100644
--- a/extensions/dns/rest/dns.rb
+++ b/extensions/dns/rest/dns.rb
@@ -3,45 +3,18 @@
 # Browser Exploitation Framework (BeEF) - http://beefproject.com
 # See the file 'doc/COPYING' for copying permission
 #
-
-# GET:
-#   * Rule count
-#   * List of rules
-
-# POST:
-#   * Add rule
-#   * Remove rule
-
-
-# /api/dns/rules
-# {
-#   "rules": [
-#       {
-#           "id": 1,
-#           "pattern": "foobar.com",
-#           "type": "Resolv::DNS::Resource::IN::A"
-#           "block": "proc {|t| ...do shit... }"
-#       },
-#
-#       {
-#       },
-#
-#       {
-#       },
-#   ]
-# }
-
-
-
 module BeEF
 module Extension
 module DNS
 
+  # This class handles the routing of RESTful API requests that query BeEF's DNS server
   class DNSRest < BeEF::Core::Router::Router
 
+    # Filters out bad requests before performing any routing
     before do
       config = BeEF::Core::Configuration.instance
 
+      # Require a valid API token from a valid IP address
       error 401 unless params[:token] == config.get('beef.api_token')
       halt  401 unless BeEF::Core::Rest.permitted_source?(request.ip)
 
@@ -53,8 +26,17 @@ module DNS
 
     # Returns the entire current DNS ruleset
     get '/rules' do
-      result = {}
-      result[:rules] = BeEF::Extension::DNS::DNS.instance.get_rules
+      result[:rules] = BeEF::Extension::DNS::DNS.instance.get_ruleset
+      result.to_json
+    end
+
+    # Returns a specific rule given its id
+    get '/rule/:id' do
+      id = params[:id]
+
+      halt 401 unless BeEF::Filters.nums_only?(id)
+
+      result = BeEF::Extension::DNS::DNS.instance.get_rule(id)
       result.to_json
     end
 
diff --git a/extensions/dns/ruby/rubydns.rb b/extensions/dns/ruby/rubydns.rb
index 222c6aec8..0af5744c0 100644
--- a/extensions/dns/ruby/rubydns.rb
+++ b/extensions/dns/ruby/rubydns.rb
@@ -84,6 +84,24 @@ module RubyDNS
       end
     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
+        result[:block] = rule.block
+      rescue DataMapper::ObjectNotFoundError => e
+        @logger.error(e.message)
+      end
+
+      result
+    end
+
   end
 
   class Transaction