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

3
require 'rex/socket'
4

5
module Rex
6
module Proto
7
module DNS
8

9
class Server
10

11
  class MockDnsClient
12
    extend Forwardable
13
    attr_reader :peerhost, :peerport, :srvsock
14

15
    def_delegators :@srvsock, :localhost, :localport, :sendto
16

17
    #
18
    # Create mock DNS client
19
    #
20
    # @param host [String] PeerHost IP address
21
    # @param port [Fixnum] PeerPort integer
22
    def initialize(host, port, sock)
23
      @peerhost = host
24
      @peerport = port
25
      @srvsock = sock
26
    end
27

28
    #
29
    # Test method to prevent GC/ObjectSpace abuse via class lookups
30
    #
31
    def mock_dns_client?
32
      true
33
    end
34

35
    def write(data)
36
      srvsock.sendto(data, peerhost, peerport)
37
    end
38
  end
39

40
  include Rex::IO::GramServer
41

42
  Packet = Rex::Proto::DNS::Packet
43
  #
44
  # Create DNS Server
45
  #
46
  # @param lhost [String] Listener address
47
  # @param lport [Fixnum] Listener port
48
  # @param udp [TrueClass, FalseClass] Listen on UDP socket
49
  # @param tcp [TrueClass, FalseClass] Listen on TCP socket
50
  # @param res [Rex::Proto::DNS::Resolver] Resolver to use, nil to create a fresh one
51
  # @param ctx [Hash] Framework context for sockets
52
  # @param dblock [Proc] Handler for :dispatch_request flow control interception
53
  # @param sblock [Proc] Handler for :send_response flow control interception
54
  #
55
  # @return [Rex::Proto::DNS::Server] DNS Server object
56
  attr_accessor :serve_tcp, :serve_udp, :fwd_res, :cache, :start_cache
57
  attr_reader :serve_udp, :serve_tcp, :sock_options, :lock, :udp_sock, :tcp_sock
58
  def initialize(lhost = '0.0.0.0', lport = 53, udp = true, tcp = false, start_cache = true, res = nil, comm = nil, ctx = {}, dblock = nil, sblock = nil)
59

60
    @serve_udp = udp
61
    @serve_tcp = tcp
62
    @sock_options = {
63
      'LocalHost' => lhost,
64
      'LocalPort' => lport,
65
      'Context'   => ctx,
66
      'Comm'      => comm
67
    }
68
    self.fwd_res = res.nil? ? Rex::Proto::DNS::Resolver.new(:comm => comm, :context => ctx) : res
69
    self.listener_thread = nil
70
    self.dispatch_request_proc = dblock
71
    self.send_response_proc = sblock
72
    self.start_cache = start_cache
73
    self.cache = Rex::Proto::DNS::Cache.new
74
    @lock = Mutex.new
75
  end
76

77
  #
78
  # Switch DNS forwarders in resolver with thread safety
79
  #
80
  # @param ns [Array, String] List of (or single) nameservers to use
81
  def switchns(ns = [])
82
    if ns.respond_to?(:split)
83
      ns = [ns]
84
    end
85
    self.lock.synchronize do
86
      self.fwd_res.nameserver = ns
87
    end
88
  end
89

90
  #
91
  # Check if server is running
92
  #
93
  def running?
94
    self.listener_thread and self.listener_thread.alive?
95
  end
96

97
  #
98
  # Start the DNS server and cache
99
  # @param start_cache [TrueClass, FalseClass] stop the cache
100
  def start
101

102
    if self.serve_udp
103
      @udp_sock = Rex::Socket::Udp.create(self.sock_options)
104
      self.listener_thread = Rex::ThreadFactory.spawn("UDPDNSServerListener", false) {
105
        monitor_listener
106
      }
107
    end
108

109
    if self.serve_tcp
110
      @tcp_sock = Rex::Socket::TcpServer.create(self.sock_options)
111
      self.tcp_sock.on_client_data_proc = Proc.new { |cli|
112
        on_client_data(cli)
113
      }
114
      self.tcp_sock.start
115
      if !self.serve_udp
116
        self.listener_thread = tcp_sock.listener_thread
117
      end
118
    end
119

120
    self.cache.start if self.start_cache
121
  end
122

123
  #
124
  # Stop the DNS server and cache
125
  #
126
  # @param flush_cache [TrueClass,FalseClass] Flush eDNS cache on stop
127
  def stop(flush_cache = false)
128
    ensure_close = [self.udp_sock, self.tcp_sock].compact
129
    begin
130
      self.listener_thread.kill if self.listener_thread.respond_to?(:kill)
131
      self.listener_thread = nil
132
    ensure
133
      while csock = ensure_close.shift
134
        csock.stop if csock.respond_to?(:stop)
135
        csock.close unless csock.respond_to?(:close) and csock.closed?
136
      end
137
    end
138
    self.cache.stop(flush_cache)
139
  end
140

141
  #
142
  # Process client request, handled with dispatch_request_proc if set
143
  #
144
  # @param cli [Rex::Socket::Tcp, Rex::Socket::Udp] Client sending the request
145
  # @param data [String] raw DNS request data
146
  def dispatch_request(cli, data)
147
    if self.dispatch_request_proc
148
      self.dispatch_request_proc.call(cli,data)
149
    else
150
      default_dispatch_request(cli,data)
151
    end
152
  end
153

154
  #
155
  # Default DNS request dispatcher, attempts to find
156
  # response records in cache or forwards request upstream
157
  #
158
  # @param cli [Rex::Socket::Tcp, Rex::Socket::Udp] Client sending the request
159
  # @param data [String] raw DNS request data
160
  def default_dispatch_request(cli,data)
161
    return if data.strip.empty?
162
    req = Packet.encode_drb(data)
163
    forward = req.dup
164
    # Find cached items, remove request from forwarded packet
165
    req.question.each do |ques|
166
      cached = self.cache.find(ques.qname, ques.qtype)
167
      if cached.empty?
168
        next
169
      else
170
        req.instance_variable_set(:@answer, (req.answer + cached).uniq)
171
        forward.question.delete(ques)
172
      end
173
    end
174
    # Forward remaining requests, cache responses
175
    if forward.question.count > 0 and @fwd_res
176
      forwarded = self.fwd_res.send(forward)
177
      req.instance_variable_set(:@answer, (req.answer + forwarded.answer).uniq)
178
      forwarded.answer.each do |ans|
179
        self.cache.cache_record(ans)
180
      end
181
      req.header.ra = true # Set recursion bit
182
    end
183
    # Finalize answers in response
184
    # Check for empty response prior to sending
185
    if req.answer.size < 1
186
      req.header.rCode = Dnsruby::RCode::NOERROR
187
    end
188
    req.header.qr = true # Set response bit
189
    send_response(cli, req.data)
190
  end
191

192
  #
193
  # Returns the hardcore alias for the DNS service
194
  #
195
  def self.hardcore_alias(*args)
196
    "#{(args[0] || '')}-#{(args[1] || '')}-#{args[5] || ''}"
197
  end
198

199
  #
200
  # DNS server.
201
  #
202
  def alias
203
    "DNS Server"
204
  end
205

206

207
protected
208
  #
209
  # This method monitors the listener socket for new connections and calls
210
  # the +on_client_connect+ callback routine.
211
  #
212
  def monitor_listener
213
    while true
214
      rds = [self.udp_sock]
215
      wds = []
216
      eds = [self.udp_sock]
217

218
      r,_,_ = ::IO.select(rds,wds,eds,1)
219

220
      if (r != nil and r[0] == self.udp_sock)
221
        buf,host,port = self.udp_sock.recvfrom(65535)
222
        # Mock up a client object for sending back data
223
        cli = MockDnsClient.new(host, port, r[0])
224
        dispatch_request(cli, buf)
225
      end
226
    end
227
  end
228

229
  #
230
  # Processes request coming from client
231
  #
232
  # @param cli [Rex::Socket::Tcp] Client sending request
233
  def on_client_data(cli)
234
    begin
235
      data = cli.read(65535)
236

237
      raise ::EOFError if not data
238
      raise ::EOFError if data.empty?
239
      dispatch_request(cli, data)
240
    rescue EOFError => e
241
      self.tcp_socket.close_client(cli) if cli
242
      raise e
243
    end
244
  end
245

246
end
247

248
end
249
end
250
end
251

252