From 2f5133e11a4573332c071325355e036f1feb2a63 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Wed, 5 Jun 2013 15:56:33 -0400
Subject: [PATCH] Changed GET handlers to return recently fixed rule data. Also
 wrapped all handlers in a begin/end block that catches internal StandardError
 exceptions.

---
 extensions/dns/dns.rb          |  2 +-
 extensions/dns/rest/dns.rb     | 17 ++++++++++++++---
 extensions/dns/ruby/rubydns.rb | 34 ++++++++++++++++++++++++++++++----
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index 78d7259db..b64b86190 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -93,7 +93,7 @@ module Dns
     # * <code>:id</code>
     # * <code>:pattern</code>
     # * <code>:type</code>
-    # * <code>:block</code>
+    # * <code>:response</code>
     #
     # @return [Array<Hash>] DNS ruleset (empty if no rules are currently loaded)
     def get_ruleset
diff --git a/extensions/dns/rest/dns.rb b/extensions/dns/rest/dns.rb
index f83568788..5fde2d05f 100644
--- a/extensions/dns/rest/dns.rb
+++ b/extensions/dns/rest/dns.rb
@@ -26,9 +26,14 @@ module Dns
 
     # Returns the entire current DNS ruleset
     get '/rules' do
-      result = {}
-      result[:rules] = BeEF::Extension::Dns::Server.instance.get_ruleset
-      result.to_json
+      begin
+        result = {}
+        result[:rules] = BeEF::Extension::Dns::Server.instance.get_ruleset
+        result.to_json
+      rescue StandardError => e
+        print_error "Internal error while retrieving DNS ruleset (#{e.message})"
+        halt 500
+      end
     end
 
     # Returns a specific rule given its id
@@ -45,6 +50,9 @@ module Dns
       rescue InvalidJsonError => e
         print_error e.message
         halt 400
+      rescue StandardError => e
+        print_error "Internal error while retrieving DNS rule with id #{id} (#{e.message})"
+        halt 500
       end
     end
 
@@ -118,6 +126,9 @@ module Dns
       rescue InvalidJsonError => e
         print_error e.message
         halt 400
+      rescue StandardError => e
+        print_error "Internal error while removing DNS rule with id #{id} (#{e.message})"
+        halt 500
       end
     end
 
diff --git a/extensions/dns/ruby/rubydns.rb b/extensions/dns/ruby/rubydns.rb
index dc3c02047..5c98dc785 100644
--- a/extensions/dns/ruby/rubydns.rb
+++ b/extensions/dns/ruby/rubydns.rb
@@ -140,8 +140,8 @@ module RubyDNS
 
         element[:id] = rule.id
         element[:pattern] = rule.pattern
-        element[:type] = rule.type
-        element[:block] = rule.block
+        element[:type] = rule.type.to_s.split('::')[-1]
+        element[:response] = parse_response(rule.block)
 
         result << element
       end
@@ -158,8 +158,8 @@ module RubyDNS
 
         result[:id] = rule.id
         result[:pattern] = rule.pattern
-        result[:type] = rule.type
-        result[:block] = rule.block
+        result[:type] = rule.type.to_s.split('::')[-1]
+        result[:response] = parse_response(rule.block)
       rescue DataMapper::ObjectNotFoundError => e
         @logger.error(e.message)
       end
@@ -183,6 +183,32 @@ module RubyDNS
       id
     end
 
+    # New method that parses response callback and returns RDATA as an array
+    def parse_response(block)
+      # Extract response arguments into an array
+      args = /(?<=respond!\().*(?=\))/.match(block).to_s.split(/,\s*/)
+
+      result = []
+
+      # Determine whether each argument is a domain name, integer, or IP address
+      args.each do |elem|
+        arg = nil
+
+        if /Name\.create\((.*)\)/.match(elem)
+          arg = $1
+        else
+          int_test = elem.to_i
+          arg = (int_test != 0 ? int_test : elem)
+        end
+
+        arg.gsub!('"', '') unless arg.is_a?(Integer)
+
+        result << arg
+      end
+
+      result
+    end
+
   end
 
   class Transaction