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

9
class Server
10

11
  class MockDnsClient
12
    attr_reader :peerhost, :peerport, :srvsock
13

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

25
    #
26
    # Test method to prevent GC/ObjectSpace abuse via class lookups
27
    #
28
    def mock_dns_client?
29
      true
30
    end
31

32
    def write(data)
33
      srvsock.sendto(data, peerhost, peerport)
34
    end
35
  end
36

37
  include Rex::IO::GramServer
38

39
  Packet = Rex::Proto::DNS::Packet
40
  #
41
  # Create DNS Server
42
  #
43
  # @param lhost [String] Listener address
44
  # @param lport [Fixnum] Listener port
45
  # @param udp [TrueClass, FalseClass] Listen on UDP socket
46
  # @param tcp [TrueClass, FalseClass] Listen on TCP socket
47
  # @param res [Rex::Proto::DNS::Resolver] Resolver to use, nil to create a fresh one
48
  # @param ctx [Hash] Framework context for sockets
49
  # @param dblock [Proc] Handler for :dispatch_request flow control interception
50
  # @param sblock [Proc] Handler for :send_response flow control interception
51
  #
52
  # @return [Rex::Proto::DNS::Server] DNS Server object
53
  attr_accessor :serve_tcp, :serve_udp, :fwd_res, :cache, :start_cache
54
  attr_reader :serve_udp, :serve_tcp, :sock_options, :lock, :udp_sock, :tcp_sock
55
  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)
56

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

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

87
  #
88
  # Check if server is running
89
  #
90
  def running?
91
    self.listener_thread and self.listener_thread.alive?
92
  end
93

94
  #
95
  # Start the DNS server and cache
96
  # @param start_cache [TrueClass, FalseClass] stop the cache
97
  def start
98

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

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

117
    self.cache.start if self.start_cache
118
  end
119

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

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

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

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

196
  #
197
  # DNS server.
198
  #
199
  def alias
200
    "DNS Server"
201
  end
202

203

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

215
      r,_,_ = ::IO.select(rds,wds,eds,1)
216

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

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

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

243
end
244

245
end
246
end
247
end
248

249