From d45bff3a5914f2117befdbc09fa8985faeea1860 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Mon, 10 Jun 2013 18:11:36 -0400
Subject: [PATCH] Improved #check_dns_response to use config file for
 address/port.

---
 test/integration/tc_dns_rest.rb | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/test/integration/tc_dns_rest.rb b/test/integration/tc_dns_rest.rb
index 9919cd6ea..b6921ebf5 100644
--- a/test/integration/tc_dns_rest.rb
+++ b/test/integration/tc_dns_rest.rb
@@ -22,6 +22,20 @@ class TC_DnsRest < Test::Unit::TestCase
 
       result  = JSON.parse(response.body)
       @@token = result['token']
+
+      $root_dir = '../../'
+      $:.unshift($root_dir)
+
+      require 'core/loader'
+
+      BeEF::Core::Configuration.new(File.join($root_dir, 'config.yaml'))
+      BeEF::Core::Configuration.instance.load_extensions_config
+
+      @@config = BeEF::Core::Configuration.instance
+    end
+
+    def shutdown
+      $root_dir = nil
     end
 
   end
@@ -362,11 +376,12 @@ class TC_DnsRest < Test::Unit::TestCase
     assert(result['id'])
   end
 
-  # TODO: Use BeEF::Core::Configuration to get address and port values.
-
   # Compares output of dig command against regex
   def check_dns_response(regex, type, pattern)
-    dig_output = `dig @localhost -p 5300 -t #{type} #{pattern}`
+    address = @@config.get('beef.extension.dns.address')
+    port    = @@config.get('beef.extension.dns.port')
+
+    dig_output = `dig @#{address} -p #{port} -t #{type} #{pattern}`
     assert_match(regex, dig_output)
   end