Added new method #remove_ruleset that clears the entire DNS ruleset.
Included unit tests as well.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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'
|
||||
|
||||
Reference in New Issue
Block a user