Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
# -*- coding: binary -*-
2

3
require 'net/dns/resolver'
4

5
module Rex
6
module Proto
7
module DNS
8

9
  ##
10
  # Provides Rex::Sockets compatible version of Net::DNS::Resolver
11
  # Modified to work with Dnsruby::Messages, their resolvers are too heavy
12
  ##
13
  class CachedResolver < Resolver
14
    include Rex::Proto::DNS::Constants
15
    attr_accessor :cache
16

17
    #
18
    # Initialize resolver with cache
19
    #
20
    # @param config [Hash] Resolver config
21
    #
22
    # @return [nil]
23
    def initialize(config = {})
24
      dns_cache_no_start = config.delete(:dns_cache_no_start)
25
      super(config)
26
      self.cache = Rex::Proto::DNS::Cache.new
27
      self.cache.start unless dns_cache_no_start
28
      return
29
    end
30

31
    #
32
    # Attempt to find answers to query in DNS cache; failing that,
33
    # send remainder of DNS queries over appropriate transport and
34
    # cache answers before returning to caller.
35
    #
36
    # @param argument [Object] An object holding the DNS message to be processed.
37
    # @param type [Fixnum] Type of record to look up
38
    # @param cls [Fixnum] Class of question to look up
39
    #
40
    # @return [Dnsruby::Message, nil] DNS response on success, nil on failure.
41
    def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
42
      case argument
43
      when Dnsruby::Message
44
        req = argument
45
      when Net::DNS::Packet, Resolv::DNS::Message
46
        req = Rex::Proto::DNS::Packet.encode_drb(argument)
47
      else
48
        net_packet = make_query_packet(argument,type,cls)
49
        # This returns a Net::DNS::Packet. Convert to Dnsruby::Message for consistency
50
        req = Rex::Proto::DNS::Packet.encode_drb(net_packet)
51
      end
52
      resolve = req.dup
53
      # Find cached items, remove request from resolved packet
54
      req.question.each do |ques|
55
        cached = self.cache.find(ques.qname, ques.qtype.to_s)
56
        next if cached.empty?
57
        req.instance_variable_set(:@answer, (req.answer + cached).uniq)
58
        resolve.question.delete(ques)
59
      end
60
      # Resolve remaining requests, cache responses
61
      if resolve.question.count > 0
62
        resolved = super(resolve, type)
63
        req.instance_variable_set(:@answer, (req.answer + resolved.answer).uniq)
64
        resolved.answer.each do |ans|
65
          self.cache.cache_record(ans)
66
        end
67
      end
68
      # Finalize answers in response
69
      # Check for empty response prior to sending
70
      req.header.rcode = Dnsruby::RCode::NOERROR if req.answer.size < 1
71
      req.header.qr = true # Set response bit
72
      return req
73
    end
74
  end
75
end
76
end
77
end
78

79