diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb index 3aaf19863..9e3f4e9db 100644 --- a/extensions/dns/dns.rb +++ b/extensions/dns/dns.rb @@ -11,17 +11,29 @@ module DNS include Singleton - # Starts DNS server run-loop. + # @!method instance + # Returns the singleton instance. + def initialize + @server = nil + @next_id = 0 + end + + # Starts the DNS server run-loop. # # @param address [String] interface address server should run on # @param port [Integer] desired server port number def run_server(address, port) EventMachine::next_tick do RubyDNS::run_server(:listen => [[:udp, address, port]]) do - upstream = RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]]) + server = self + BeEF::Extension::DNS::DNS.instance.instance_eval { @server = server } + BeEF::Extension::DNS::DNS.instance.load_rules + # Pass unmatched queries upstream to root nameservers otherwise do |transaction| - transaction.passthrough!(upstream) + transaction.passthrough!( + RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]]) + ) end end end @@ -29,25 +41,56 @@ module DNS # Adds a new DNS rule or "resource record". Does nothing if rule is already present. # - # @param name [String] name of query - # @param type [String] query type (e.g. A, CNAME, MX, NS, etc.) - # @param value [String] response to send back to resolver - def add_rule(name, type, value) - catch(:match) do - BeEF::Core::Models::DNS.each do |rule| - n = rule.name - t = rule.type - v = rule.value + # @example Adds an A record for foobar.com with the value 1.2.3.4 + # + # dns = BeEF::Extension::DNS::DNS.instance + # + # id = dns.add_rule('foobar.com', Resolv::DNS::Resource::IN::A) do |transaction| + # transaction.respond!('1.2.3.4') + # end + # + # @param pattern [String, Regexp] query pattern to recognize + # @param type [Resolv::DNS::Resource::IN] resource record type (e.g. A, CNAME, NS, etc.) + # + # @yieldparam [RubyDNS::Transaction] details of query question and response + # + # @return [Integer] unique id for use with {#remove_rule} + # + # @see #remove_rule + # @see http://rubydoc.info/gems/rubydns/RubyDNS/Transaction + def add_rule(pattern, type, &block) + @next_id += 1 + @server.match(@next_id, pattern, type, block) + @next_id + end - throw :match if [n, t, v] == [name, type, value] - end + # Removes the given DNS rule. Any future queries for it will be passed through. + # + # @param id [Integer] id returned from {#add_rule} + # @see #add_rule + def remove_rule(id) + @server.remove_rule(id) + @next_id -= 1 + end - BeEF::Core::Models::DNS.create( - :name => name, - :type => type, - :value => value - ) - end + # Loads all rules from the database at server startup + def load_rules + # TODO Load rules from database + end + + private + + # Convenience method for fully-qualifying Resolv::DNS::Resource::IN types + def parse_type(type) + eval "Resolv::DNS::Resource::IN::#{type}" + end + + # Convenience method for generating proper server responses + def parse_response(type, value) + response = 'value.to_s' + response = 'Resolv::DNS::Name.create(value)' if type =~ /(CNAME|NS|PTR)/ + + eval response end end diff --git a/extensions/dns/ruby/rubydns.rb b/extensions/dns/ruby/rubydns.rb index 39ab1af93..b7a4bdf4c 100644 --- a/extensions/dns/ruby/rubydns.rb +++ b/extensions/dns/ruby/rubydns.rb @@ -9,8 +9,7 @@ require 'rubydns' module RubyDNS - # Behaves exactly the same, except without any output and an added periodic - # timer that checks for new DNS rules every five seconds + # Behaves exactly the same, except without any logger output def self.run_server(options = {}, &block) server = RubyDNS::Server.new(&block) @@ -27,15 +26,41 @@ module RubyDNS end end - server.load_rules - EventMachine.add_periodic_timer(5) { server.check_rules } - server.fire(:start) end server.fire(:stop) end + class Server + + attr_accessor :rules + + class Rule + + attr_accessor :id + + # Now uses an 'id' parameter to uniquely identify rules + def initialize(id, pattern, callback) + @id = id + @pattern = pattern + @callback = callback + end + + end + + # Now uses an 'id' parameter to uniquely identify rules + def match(id, *pattern, block) + @rules << Rule.new(id, pattern, block) + end + + # New method that removes a rule given its id + def remove_rule(id) + @rules.delete_if { |rule| rule.id == id } + end + + end + class Transaction # Behaves exactly the same, except using debug logger instead of info @@ -56,53 +81,4 @@ module RubyDNS end - class Server - - # Reads current DNS entries in database and adds them as new rules - def load_rules - rules = get_rules - @rule_count = rules.count - - rules.each do |rule| - match(rule[0], parse_type(rule[1])) do |transaction| - transaction.respond!(rule[2]) - end - end - end - - # Re-loads ruleset if new entries have been added to database - def check_rules - load_rules if get_rules.count != @rule_count - end - - private - - # Returns an AoA where each element is a rule of the form [name, type, value] - def get_rules - rules = [] - - BeEF::Core::Models::DNS.each do |record| - name = record.name - type = record.type - value = record.value - - rules << [name, type, value] - end - - rules - end - - # Convenience method for fully-qualifying Resolv::DNS::Resource types - def parse_type(type) - resolv = 'Resolv::DNS::Resource' - - if type =~ /(A|AAAA|SRV|WKS)/ - resolv += '::IN' - end - - eval "#{resolv}::#{type}" - end - - end - end