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

3
require 'rex/socket'
4

5
module Rex
6
module Proto
7
module DNS
8
  class Cache
9
    attr_reader :records, :lock, :monitor_thread
10
    include Rex::Proto::DNS::Constants
11
    # class DNSRecordError < ::Exception
12
    #
13
    # Create DNS cache
14
    #
15
    def initialize
16
      @records = {}
17
      @lock = Mutex.new
18
    end
19

20
    #
21
    # Find entries in cache, substituting names for '*' in return
22
    #
23
    # @param search [String] Name or address to search for
24
    # @param type [Dnsruby::Types] Record type to search for
25
    #
26
    # @return [Array] Records found
27
    def find(search, type = Dnsruby::Types::A)
28
      self.records.select do |record,expire|
29
        record.type == type and (expire < 1 or expire > ::Time.now.to_i) and
30
        (
31
          record.name == '*' or
32
          record.name.to_s == search.to_s or record.name.to_s[0..-2] == search.to_s or
33
          ( record.respond_to?(:address) and record.address.to_s == search.to_s )
34
        )
35
      end.keys.map do |record|
36
        if search.to_s.match(MATCH_HOSTNAME) and record.name == '*'
37
          record = Dnsruby::RR.create(name: name, type: type, address: address)
38
        else
39
          record
40
        end
41
      end
42
    end
43

44
    #
45
    # Add record to cache, only when "running"
46
    #
47
    # @param record [Dnsruby::RR] Record to cache
48
    def cache_record(record, expire: true)
49
      return unless @monitor_thread
50

51
      unless record.is_a?(Dnsruby::RR)
52
        raise "Invalid record for cache entry (not an RR) - #{record.inspect}"
53
      end
54

55
      unless (!record.respond_to?(:address) || Rex::Socket.is_ip_addr?(record.address.to_s))
56
        raise "Invalid record for cache entry (no IP address) - #{record.inspect}"
57
      end
58

59
      unless record.name.to_s.match(MATCH_HOSTNAME)
60
        raise "Invalid record for cache entry (invalid hostname) - #{record.inspect}"
61
      end
62

63
      add(record, expire ? (::Time.now.to_i + record.ttl) : 0)
64
    end
65

66
    #
67
    # Delete all cache entries, this is different from pruning because the
68
    # record's expiration is ignored
69
    #
70
    def flush
71
      self.records.each {|rec, _| delete(rec)}
72
    end
73

74
    #
75
    # Prune cache entries
76
    #
77
    # @param before [Fixnum] Time in seconds before which records are evicted
78
    def prune(before = ::Time.now.to_i)
79
      self.records.select do |rec, expire|
80
        expire > 0 and expire < before
81
      end.each {|rec, exp| delete(rec)}
82
    end
83

84
    #
85
    # Start the cache monitor
86
    #
87
    def start
88
      @monitor_thread = Rex::ThreadFactory.spawn("DNSServerCacheMonitor", false) {
89
        while true
90
          prune
91
          Rex::ThreadSafe.sleep(0.5)
92
        end
93
      } unless @monitor_thread
94
    end
95

96
    #
97
    # Stop the cache monitor
98
    #
99
    # @param flush [TrueClass,FalseClass] Remove non-static entries
100
    def stop(flush = false)
101
      self.monitor_thread.kill unless @monitor_thread.nil?
102
      @monitor_thread = nil
103
      if flush
104
        self.records.select do |rec, expire|
105
          rec.ttl > 0
106
        end.each {|rec| delete(rec)}
107
      end
108
    end
109

110
    protected
111

112
    #
113
    # Add a record to the cache with thread safety
114
    #
115
    # @param record [Dnsruby::RR] Record to add
116
    # @param expire [Fixnum] Time in seconds when record becomes stale
117
    def add(record, expire = 0)
118
      self.lock.synchronize do
119
        self.records[record] = expire
120
      end
121
    end
122

123
    #
124
    # Delete a record from the cache with thread safety
125
    #
126
    # @param record [Dnsruby::RR] Record to delete
127
    def delete(record)
128
      self.lock.synchronize do
129
        self.records.delete(record)
130
      end
131
    end
132
  end # Cache
133
end
134
end
135
end
136

137