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 '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)
49
      return unless @monitor_thread
50
      if record.is_a?(Dnsruby::RR) and
51
      (!record.respond_to?(:address) or Rex::Socket.is_ip_addr?(record.address.to_s)) and
52
      record.name.to_s.match(MATCH_HOSTNAME)
53
        add(record, ::Time.now.to_i + record.ttl)
54
      else
55
        raise "Invalid record for cache entry - #{record.inspect}"
56
      end
57
    end
58

59
    #
60
    # Delete all cache entries, this is different from pruning because the
61
    # record's expiration is ignored
62
    #
63
    def flush
64
      self.records.each {|rec, _| delete(rec)}
65
    end
66

67
    #
68
    # Prune cache entries
69
    #
70
    # @param before [Fixnum] Time in seconds before which records are evicted
71
    def prune(before = ::Time.now.to_i)
72
      self.records.select do |rec, expire|
73
        expire > 0 and expire < before
74
      end.each {|rec, exp| delete(rec)}
75
    end
76

77
    #
78
    # Start the cache monitor
79
    #
80
    def start
81
      @monitor_thread = Rex::ThreadFactory.spawn("DNSServerCacheMonitor", false) {
82
        while true
83
          prune
84
          Rex::ThreadSafe.sleep(0.5)
85
        end
86
      } unless @monitor_thread
87
    end
88

89
    #
90
    # Stop the cache monitor
91
    #
92
    # @param flush [TrueClass,FalseClass] Remove non-static entries
93
    def stop(flush = false)
94
      self.monitor_thread.kill unless @monitor_thread.nil?
95
      @monitor_thread = nil
96
      if flush
97
        self.records.select do |rec, expire|
98
          rec.ttl > 0
99
        end.each {|rec| delete(rec)}
100
      end
101
    end
102

103
    protected
104

105
    #
106
    # Add a record to the cache with thread safety
107
    #
108
    # @param record [Dnsruby::RR] Record to add
109
    # @param expire [Fixnum] Time in seconds when record becomes stale
110
    def add(record, expire = 0)
111
      self.lock.synchronize do
112
        self.records[record] = expire
113
      end
114
    end
115

116
    #
117
    # Delete a record from the cache with thread safety
118
    #
119
    # @param record [Dnsruby::RR] Record to delete
120
    def delete(record)
121
      self.lock.synchronize do
122
        self.records.delete(record)
123
      end
124
    end
125
  end # Cache
126
end
127
end
128
end
129

130