Added new method #remove_ruleset that clears the entire DNS ruleset.

Included unit tests as well.
This commit is contained in:
soh_cah_toa
2013-07-17 18:16:46 -04:00
parent 9cfb98963d
commit 95d0ddbe87
3 changed files with 53 additions and 28 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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'