1#
2# This file is a part of DNSViz, a tool suite for DNS/DNSSEC monitoring,
3# analysis, and visualization.
4# Created by Casey Deccio (casey@deccio.net)
5#
6# Copyright 2012-2014 Sandia Corporation. Under the terms of Contract
7# DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
8# certain rights in this software.
9#
10# Copyright 2014-2016 VeriSign, Inc.
11#
12# Copyright 2016-2021 Casey Deccio
13#
14# DNSViz is free software; you can redistribute it and/or modify
15# it under the terms of the GNU General Public License as published by
16# the Free Software Foundation; either version 2 of the License, or
17# (at your option) any later version.
18#
19# DNSViz is distributed in the hope that it will be useful,
20# but WITHOUT ANY WARRANTY; without even the implied warranty of
21# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
22# GNU General Public License for more details.
23#
24# You should have received a copy of the GNU General Public License along
25# with DNSViz.  If not, see <http://www.gnu.org/licenses/>.
26#
27
28from __future__ import unicode_literals
29
30import binascii
31import bisect
32import copy
33import errno
34import io
35import socket
36import struct
37import time
38
39# minimal support for python2.6
40try:
41    from collections import OrderedDict
42except ImportError:
43    from ordereddict import OrderedDict
44
45# python3/python2 dual compatibility
46try:
47    import queue
48except ImportError:
49    import Queue as queue
50
51import dns.edns, dns.exception, dns.flags, dns.message, dns.rcode, \
52        dns.rdataclass, dns.rdatatype
53
54from .ipaddr import *
55from .response import *
56from . import transport
57from .format import latin1_binary_to_string as lb2s
58
59RETRY_CAUSE_NETWORK_ERROR = RESPONSE_ERROR_NETWORK_ERROR = 1
60RETRY_CAUSE_FORMERR = RESPONSE_ERROR_FORMERR = 2
61RETRY_CAUSE_TIMEOUT = RESPONSE_ERROR_TIMEOUT = 3
62RETRY_CAUSE_OTHER = RESPONSE_ERROR_OTHER = 4
63RETRY_CAUSE_TC_SET = 5
64RETRY_CAUSE_RCODE = RESPONSE_ERROR_INVALID_RCODE = 6
65RETRY_CAUSE_DIAGNOSTIC = 7
66retry_causes = {
67        RETRY_CAUSE_NETWORK_ERROR: 'NETWORK_ERROR',
68        RETRY_CAUSE_FORMERR: 'FORMERR',
69        RETRY_CAUSE_TIMEOUT: 'TIMEOUT',
70        RETRY_CAUSE_OTHER: 'ERROR',
71        RETRY_CAUSE_TC_SET: 'TC',
72        RETRY_CAUSE_RCODE: 'INVALID_RCODE',
73        RETRY_CAUSE_DIAGNOSTIC: 'DIAGNOSTIC'
74}
75retry_cause_codes = {
76        'NETWORK_ERROR': RETRY_CAUSE_NETWORK_ERROR,
77        'FORMERR': RETRY_CAUSE_FORMERR,
78        'TIMEOUT': RETRY_CAUSE_TIMEOUT,
79        'ERROR': RETRY_CAUSE_OTHER,
80        'TC': RETRY_CAUSE_TC_SET,
81        'INVALID_RCODE': RETRY_CAUSE_RCODE,
82        'DIAGNOSTIC': RETRY_CAUSE_DIAGNOSTIC,
83}
84response_errors = {
85        RESPONSE_ERROR_NETWORK_ERROR: retry_causes[RETRY_CAUSE_NETWORK_ERROR],
86        RESPONSE_ERROR_FORMERR: retry_causes[RETRY_CAUSE_FORMERR],
87        RESPONSE_ERROR_TIMEOUT: retry_causes[RETRY_CAUSE_TIMEOUT],
88        RESPONSE_ERROR_OTHER: retry_causes[RETRY_CAUSE_OTHER],
89        RESPONSE_ERROR_INVALID_RCODE: retry_causes[RETRY_CAUSE_RCODE]
90}
91response_error_codes = {
92        retry_causes[RETRY_CAUSE_NETWORK_ERROR]: RESPONSE_ERROR_NETWORK_ERROR,
93        retry_causes[RETRY_CAUSE_FORMERR]: RESPONSE_ERROR_FORMERR,
94        retry_causes[RETRY_CAUSE_TIMEOUT]: RESPONSE_ERROR_TIMEOUT,
95        retry_causes[RETRY_CAUSE_OTHER]: RESPONSE_ERROR_OTHER,
96        retry_causes[RETRY_CAUSE_RCODE]: RESPONSE_ERROR_INVALID_RCODE
97}
98
99RETRY_ACTION_NO_CHANGE = 1
100RETRY_ACTION_USE_TCP = 2
101RETRY_ACTION_USE_UDP = 3
102RETRY_ACTION_SET_FLAG = 4
103RETRY_ACTION_CLEAR_FLAG = 5
104RETRY_ACTION_DISABLE_EDNS = 6
105RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD = 7
106RETRY_ACTION_SET_EDNS_FLAG = 8
107RETRY_ACTION_CLEAR_EDNS_FLAG = 9
108RETRY_ACTION_ADD_EDNS_OPTION = 10
109RETRY_ACTION_REMOVE_EDNS_OPTION = 11
110RETRY_ACTION_CHANGE_SPORT = 12
111RETRY_ACTION_CHANGE_EDNS_VERSION = 13
112RETRY_ACTION_UPDATE_DNS_COOKIE = 14
113retry_actions = {
114        RETRY_ACTION_NO_CHANGE: 'NO_CHANGE',
115        RETRY_ACTION_USE_TCP: 'USE_TCP', # implies CHANGE_SPORT
116        RETRY_ACTION_USE_UDP: 'USE_UDP', # implies CHANGE_SPORT
117        RETRY_ACTION_SET_FLAG: 'SET_FLAG', # implies CHANGE_SPORT
118        RETRY_ACTION_CLEAR_FLAG: 'CLEAR_FLAG', # implies CHANGE_SPORT
119        RETRY_ACTION_DISABLE_EDNS: 'DISABLE_EDNS', # implies CHANGE_SPORT
120        RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD: 'CHANGE_UDP_MAX_PAYLOAD', # implies USE_UDP, CHANGE_SPORT
121        RETRY_ACTION_SET_EDNS_FLAG: 'SET_EDNS_FLAG', # implies CHANGE_SPORT
122        RETRY_ACTION_CLEAR_EDNS_FLAG: 'CLEAR_EDNS_FLAG', # implies CHANGE_SPORT
123        RETRY_ACTION_ADD_EDNS_OPTION: 'ADD_EDNS_OPTION', # implies CHANGE_SPORT
124        RETRY_ACTION_REMOVE_EDNS_OPTION: 'REMOVE_EDNS_OPTION', # implies CHANGE_SPORT
125        RETRY_ACTION_CHANGE_SPORT: 'CHANGE_SPORT',
126        RETRY_ACTION_CHANGE_EDNS_VERSION: 'CHANGE_EDNS_VERSION', # implies CHANGE_SPORT
127        RETRY_ACTION_UPDATE_DNS_COOKIE: 'UPDATE_DNS_COOKIE', # implies CHANGE_SPORT
128}
129retry_action_codes = {
130        'NO_CHANGE': RETRY_ACTION_NO_CHANGE,
131        'USE_TCP': RETRY_ACTION_USE_TCP,
132        'USE_UDP': RETRY_ACTION_USE_UDP,
133        'SET_FLAG': RETRY_ACTION_SET_FLAG,
134        'CLEAR_FLAG': RETRY_ACTION_CLEAR_FLAG,
135        'DISABLE_EDNS': RETRY_ACTION_DISABLE_EDNS,
136        'CHANGE_UDP_MAX_PAYLOAD': RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD,
137        'SET_EDNS_FLAG': RETRY_ACTION_SET_EDNS_FLAG,
138        'CLEAR_EDNS_FLAG': RETRY_ACTION_CLEAR_EDNS_FLAG,
139        'ADD_EDNS_OPTION': RETRY_ACTION_ADD_EDNS_OPTION,
140        'REMOVE_EDNS_OPTION': RETRY_ACTION_REMOVE_EDNS_OPTION,
141        'CHANGE_SPORT': RETRY_ACTION_CHANGE_SPORT,
142        'CHANGE_EDNS_VERSION': RETRY_ACTION_CHANGE_EDNS_VERSION,
143        'UPDATE_DNS_COOKIE': RETRY_ACTION_UPDATE_DNS_COOKIE,
144}
145
146DNS_COOKIE_NO_COOKIE = 0
147DNS_COOKIE_CLIENT_COOKIE_ONLY = 1
148DNS_COOKIE_SERVER_COOKIE_FRESH = 2
149DNS_COOKIE_SERVER_COOKIE_STATIC = 3
150DNS_COOKIE_SERVER_COOKIE_BAD = 4
151DNS_COOKIE_IMPROPER_LENGTH = 5
152
153MIN_QUERY_TIMEOUT = 0.1
154MAX_CNAME_REDIRECTION = 40
155
156class AcceptResponse(Exception):
157    '''An exception raised to stop the process of retrying DNS queries when an
158    acceptable response or error condition has been satisfied.'''
159    pass
160
161class BindError(Exception):
162    '''An error resulting from unsuccessfully trying to bind to an address or port.'''
163    pass
164
165class SourceAddressBindError(BindError):
166    '''An error resulting from unsuccessfully trying to bind to an address.'''
167    pass
168
169class PortBindError(BindError):
170    '''An error resulting from unsuccessfully trying to bind to a port.'''
171    pass
172
173class NoValidServersToQuery(Exception):
174    '''An exception raised when a query is executed and the collective
175    transport handlers designated don't have the proper network capabilities to
176    issue queries to all the servers.'''
177    pass
178
179class DNSQueryRetryAttempt:
180    '''A failed attempt at a DNS query that invokes a subsequent retry.'''
181
182    def __init__(self, response_time, cause, cause_arg, action, action_arg):
183        self.response_time = response_time
184        self.cause = cause
185        self.cause_arg = cause_arg
186        self.action = action
187        self.action_arg = action_arg
188
189    def __repr__(self):
190        return '<Retry: %s -> %s>' % (retry_causes[self.cause], retry_actions[self.action])
191
192    def serialize(self):
193        '''Return a serialized version of the query.'''
194
195        d = OrderedDict()
196        d['time_elapsed'] = int(self.response_time * 1000)
197        d['cause'] = retry_causes.get(self.cause, 'UNKNOWN')
198        if self.cause_arg is not None:
199            if self.cause == RETRY_CAUSE_NETWORK_ERROR:
200                errno_name = errno.errorcode.get(self.cause_arg, None)
201                if errno_name is not None:
202                    d['cause_arg'] = errno_name
203            else:
204                d['cause_arg'] = self.cause_arg
205        d['action'] = retry_actions.get(self.action, 'UNKNOWN')
206        if self.action_arg is not None:
207            d['action_arg'] = self.action_arg
208        return d
209
210    @classmethod
211    def deserialize(cls, d):
212        '''Return an instance built from a serialized version of the
213        DNSQueryRetryAttempt.'''
214
215        # compatibility with version 1.0
216        if 'response_time' in d:
217            response_time = d['response_time']
218        else:
219            response_time = d['time_elapsed']/1000.0
220        cause = retry_cause_codes[d['cause']]
221        if 'cause_arg' in d:
222            if cause == RETRY_CAUSE_NETWORK_ERROR:
223                # compatibility with version 1.0
224                if isinstance(d['cause_arg'], int):
225                    cause_arg = d['cause_arg']
226                else:
227                    if hasattr(errno, d['cause_arg']):
228                        cause_arg = getattr(errno, d['cause_arg'])
229                    else:
230                        cause_arg = None
231            else:
232                cause_arg = d['cause_arg']
233        else:
234            cause_arg = None
235        action = retry_action_codes[d['action']]
236        if 'action_arg' in d:
237            action_arg = d['action_arg']
238        else:
239            action_arg = None
240        return DNSQueryRetryAttempt(response_time, cause, cause_arg, action, action_arg)
241
242class DNSResponseHandlerFactory(object):
243    '''A factory class that holds arguments to create a DNSResponseHandler instance.'''
244
245    def __init__(self, cls, *args, **kwargs):
246        self._cls = cls
247        self._args = args
248        self._kwargs = kwargs
249
250    def build(self):
251        '''Instantiate a DNSResponseHandler with the args and kwargs saved with the
252        initialization of this factory.'''
253
254        obj = self._cls.__new__(self._cls, *self._args, __instantiate=True, **self._kwargs)
255        obj.__init__(*self._args, **self._kwargs)
256        return obj
257
258class DNSResponseHandler(object):
259    '''A base class for handling DNS responses (or exceptions) arising from a
260    query attempt.'''
261
262    def __new__(cls, *args, **kwargs):
263        '''Redirect the instantiation of a DNSResponseHandler to create instead a Factory,
264        from which a DNSResponseHandler in turn is built.'''
265
266
267        if kwargs.pop('__instantiate', None):
268            return super(DNSResponseHandler, cls).__new__(cls)
269        return DNSResponseHandlerFactory(cls, *args, **kwargs)
270
271    def set_context(self, params, history, request):
272        '''Set local parameters pertaining to DNS query.'''
273
274        self._params = params
275        self._history = history
276        self._request = request
277
278    def handle(self, response_wire, response, response_time):
279        '''Handle a DNS response.  The response might be an actual DNS message or some type
280        of exception that was raised during query.'''
281
282        raise NotImplemented
283
284    def _get_retry_qty(self, cause):
285        '''Return the number of retries associated with the DNS query, optionally limited to
286        those with a given cause.'''
287
288        if cause is None:
289            return len(self._history)
290
291        total = 0
292        for i in range(len(self._history) - 1, -1, -1):
293            if self._history[i].cause == cause:
294                total += 1
295            else:
296                break
297        return total
298
299    def _get_num_timeouts(self, response):
300        '''Return the number of retries attributed to timeouts.'''
301
302        if isinstance(response, dns.exception.Timeout):
303            return self._get_retry_qty(RETRY_CAUSE_TIMEOUT) + 1
304        return 0
305
306    def _get_num_network_errors(self, response):
307        '''Return the number of retries attributed to network errors.'''
308
309        if isinstance(response, (socket.error, EOFError)):
310            return self._get_retry_qty(RETRY_CAUSE_NETWORK_ERROR) + 1
311        return 0
312
313class ActionIndependentDNSResponseHandler(DNSResponseHandler):
314    '''A DNSResponseHandler that is consulted regardless of whether or not
315    the response was "handled" previously by another handler.'''
316
317    pass
318
319class RetryOnNetworkErrorHandler(DNSResponseHandler):
320    '''Retry the query after some exponentially growing wait period upon a
321    network error.'''
322
323    def __init__(self, max_errors):
324        self._max_errors = max_errors
325
326    def handle(self, response_wire, response, response_time):
327        errors = self._get_num_network_errors(response)
328        if errors >= self._max_errors:
329            raise AcceptResponse()
330
331        if isinstance(response, (socket.error, EOFError)):
332            if hasattr(response, 'errno'):
333                errno1 = response.errno
334            else:
335                errno1 = None
336            self._params['wait'] = 0.2*(2**errors)
337
338            if self._params['tcp']:
339                action = RETRY_ACTION_CHANGE_SPORT
340            else:
341                action = RETRY_ACTION_NO_CHANGE
342            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_NETWORK_ERROR, errno1, action, None)
343
344class UseTCPOnTCFlagHandler(DNSResponseHandler):
345    '''Retry with TCP if the TC flag is set in the response.'''
346
347    def handle(self, response_wire, response, response_time):
348        # python3/python2 dual compatibility
349        if isinstance(response_wire, str):
350            map_func = lambda x: ord(x)
351        else:
352            map_func = lambda x: x
353
354        if response_wire is not None and map_func(response_wire[2]) & 0x02:
355            self._params['tcp'] = True
356            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TC_SET, len(response_wire), RETRY_ACTION_USE_TCP, None)
357
358class DisableEDNSOnFormerrHandler(DNSResponseHandler):
359    '''Disable EDNS if there was some type of issue parsing the message.  Some
360    servers don't handle EDNS appropriately.'''
361
362    def handle(self, response_wire, response, response_time):
363        if isinstance(response, (struct.error, dns.exception.FormError)) and self._request.edns >= 0:
364            self._request.use_edns(False)
365            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_FORMERR, None, RETRY_ACTION_DISABLE_EDNS, None)
366
367class ReduceUDPMaxPayloadOnTimeoutHandler(DNSResponseHandler):
368    '''Reduce the EDNS UDP max payload after a given number of timeouts.  Some
369    servers attempt to send payloads that exceed their PMTU.'''
370
371    def __init__(self, reduced_payload, timeouts):
372        self._reduced_payload = reduced_payload
373        self._timeouts = timeouts
374
375    def handle(self, response_wire, response, response_time):
376        timeouts = self._get_num_timeouts(response)
377        if not self._params['tcp'] and timeouts >= self._timeouts and self._request.payload > self._reduced_payload:
378            self._request.payload = self._reduced_payload
379            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, self._reduced_payload)
380
381class ClearEDNSFlagOnTimeoutHandler(DNSResponseHandler):
382    '''Clear an EDNS flag after a given number of timeouts.'''
383
384    def __init__(self, flag, timeouts):
385        self._flag = flag
386        self._timeouts = timeouts
387
388    def handle(self, response_wire, response, response_time):
389        timeouts = self._get_num_timeouts(response)
390        if not self._params['tcp'] and timeouts >= self._timeouts and (self._request.ednsflags & self._flag):
391            self._request.ednsflags &= ~(self._flag & 0xffff)
392            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CLEAR_EDNS_FLAG, self._flag)
393
394class ChangeEDNSVersionOnTimeoutHandler(DNSResponseHandler):
395    '''Change EDNS version after a given number of timeouts.'''
396
397    def __init__(self, edns, timeouts):
398        self._edns = edns
399        self._timeouts = timeouts
400
401    def handle(self, response_wire, response, response_time):
402        timeouts = self._get_num_timeouts(response)
403        if not self._params['tcp'] and timeouts >= self._timeouts and self._request.edns != self._edns:
404            self._request.use_edns(self._edns, self._request.ednsflags, self._request.payload, options=self._request.options)
405            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_EDNS_VERSION, self._edns)
406
407class RemoveEDNSOptionOnTimeoutHandler(DNSResponseHandler):
408    '''Remove EDNS option after a given number of timeouts.'''
409
410    def __init__(self, timeouts):
411        self._timeouts = timeouts
412
413    def handle(self, response_wire, response, response_time):
414        timeouts = self._get_num_timeouts(response)
415        try:
416            opt = self._request.options[0]
417        except IndexError:
418            opt = None
419        if not self._params['tcp'] and timeouts >= self._timeouts and opt is not None:
420            self._request.options.remove(opt)
421            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_REMOVE_EDNS_OPTION, opt.otype)
422
423class DisableEDNSOnTimeoutHandler(DNSResponseHandler):
424    '''Disable EDNS after a given number of timeouts.  Some servers don't
425    respond to EDNS queries.'''
426
427    def __init__(self, timeouts):
428        self._timeouts = timeouts
429
430    def handle(self, response_wire, response, response_time):
431        timeouts = self._get_num_timeouts(response)
432        if not self._params['tcp'] and timeouts >= self._timeouts and self._request.edns >= 0:
433            self._request.use_edns(False)
434            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_DISABLE_EDNS, None)
435
436class SetFlagOnRcodeHandler(DNSResponseHandler):
437    '''Set a flag when a given rcode is returned.  One example of the use of
438    this class is to determine if the cause of the SERVFAIL is related to DNSSEC
439    validation failure by retrying with the CD flag.'''
440
441    def __init__(self, flag, rcode):
442        self._flag = flag
443        self._rcode = rcode
444
445    def handle(self, response_wire, response, response_time):
446        if isinstance(response, dns.message.Message) and response.rcode() == self._rcode and not self._request.flags & self._flag:
447            self._request.flags |= self._flag
448            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_RCODE, self._rcode, RETRY_ACTION_SET_FLAG, self._flag)
449
450class DisableEDNSOnRcodeHandler(DNSResponseHandler):
451    '''Disable EDNS if the RCODE in the response indicates that the server
452    doesn't implement EDNS.'''
453
454    def handle(self, response_wire, response, response_time):
455        if isinstance(response, dns.message.Message) and response.rcode() in (dns.rcode.NOTIMP, dns.rcode.FORMERR, dns.rcode.SERVFAIL) and self._request.edns >= 0:
456            self._request.use_edns(False)
457            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_RCODE, response.rcode(), RETRY_ACTION_DISABLE_EDNS, None)
458
459class RemoveEDNSOptionOnRcodeHandler(DNSResponseHandler):
460    '''Remove an EDNS option if the RCODE in the response indicates that the
461    server didn't handle the request properly.'''
462
463    def __init__(self, rcode):
464        self._rcode = rcode
465
466    def handle(self, response_wire, response, response_time):
467        try:
468            opt = self._request.options[0]
469        except IndexError:
470            opt = None
471        if isinstance(response, dns.message.Message) and response.rcode() == self._rcode and opt is not None:
472            self._request.options.remove(opt)
473            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_RCODE, response.rcode(), RETRY_ACTION_REMOVE_EDNS_OPTION, opt.otype)
474
475class AddServerCookieOnBADCOOKIE(DNSResponseHandler):
476    '''Update the DNS Cookie EDNS option with the server cookie when a
477    BADCOOKIE rcode is received.'''
478
479    def _add_server_cookie(self, response):
480        try:
481            client_opt = [o for o in self._request.options if o.otype == 10][0]
482        except IndexError:
483            return False
484        try:
485            server_opt = [o for o in response.options if o.otype == 10][0]
486        except IndexError:
487            return False
488        client_cookie = client_opt.data[:8]
489        server_cookie1 = client_opt.data[8:]
490        server_cookie2 = server_opt.data[8:]
491        if server_cookie1 == server_cookie2:
492            return False
493        client_opt.data = client_cookie + server_cookie2
494        return True
495
496    def handle(self, response_wire, response, response_time):
497        if isinstance(response, dns.message.Message) and response.rcode() == 23:
498            if self._add_server_cookie(response):
499                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_RCODE, response.rcode(), RETRY_ACTION_UPDATE_DNS_COOKIE, None)
500
501class UseUDPOnTimeoutHandler(DNSResponseHandler):
502    '''Revert to UDP if TCP connectivity fails.'''
503
504    def __init__(self, timeouts):
505        self._timeouts = timeouts
506
507    def handle(self, response_wire, response, response_time):
508        timeouts = self._get_num_timeouts(response)
509        if timeouts >= self._timeouts and self._params['tcp']:
510            self._params['tcp'] = False
511            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_USE_UDP, None)
512
513class UseUDPOnNetworkErrorHandler(DNSResponseHandler):
514    '''Retry the query after some exponentially growing wait period upon a
515    network error.'''
516
517    def __init__(self, max_errors):
518        self._max_errors = max_errors
519
520    def handle(self, response_wire, response, response_time):
521        errors = self._get_num_network_errors(response)
522        if errors >= self._max_errors and self._params['tcp']:
523            if hasattr(response, 'errno'):
524                errno1 = response.errno
525            else:
526                errno1 = None
527            self._params['tcp'] = False
528            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_NETWORK_ERROR, errno1, RETRY_ACTION_USE_UDP, None)
529
530        if isinstance(response, (socket.error, EOFError)):
531            if hasattr(response, 'errno'):
532                errno1 = response.errno
533            else:
534                errno1 = None
535            self._params['wait'] = 0.2*(2**errors)
536
537            if self._params['tcp']:
538                action = RETRY_ACTION_CHANGE_SPORT
539            else:
540                action = RETRY_ACTION_NO_CHANGE
541            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_NETWORK_ERROR, errno1, action, None)
542
543class PMTUBoundingHandler(DNSResponseHandler):
544    # define states
545    START = 1
546        # if TIMEOUT -> reduce payload -> REDUCED_PAYLOAD
547        # else -> return (pass through to other handlers)
548    REDUCED_PAYLOAD = 2
549        # if TIMEOUT -> return
550        # if TC -> set lower bound, use TCP -> TCP_FOR_TRUNCATE
551        # if error -> return
552        # else -> set lower bound (msg size), use TCP -> TCP_FOR_UPPER_BOUND
553    USE_TCP = 3
554        # if TIMEOUT -> return
555        # if error -> return
556        # else -> set upper bound, set increase payload (msg payload - 1) -> TCP_MINUS_ONE
557    TCP_MINUS_ONE = 5
558        # if TIMEOUT -> reduce payload (upper - lower)/2 -> PICKLE
559        # if errors of some sort (maybe with subhandlers?) -> return
560        # else -> keep upper bound, return
561    PICKLE = 6
562        # if upper - lower <= 1 -> use TCP -> TCP_FINAL
563        # if TIMEOUT -> set upper bound, reduce payload ((upper - lower)/2 - lower)/2, PICKLE
564        # -> TC???
565        # if error -> return
566        # else -> set lower bound, increase payload (upper - (upper - lower)/2)/2, PICKLE
567    TCP_FINAL = 7
568    INVALID = 8
569
570    def __init__(self, reduced_payload, initial_timeouts, max_timeouts, bounding_timeout):
571        self._reduced_payload = reduced_payload
572        self._initial_timeouts = initial_timeouts
573        self._max_timeouts = max_timeouts
574        self._bounding_timeout = bounding_timeout
575
576        self._lower_bound = None
577        self._upper_bound = None
578        self._water_mark = None
579        self._state = self.START
580
581    def handle(self, response_wire, response, response_time):
582        if self._state == self.INVALID:
583            return
584
585        # python3/python2 dual compatibility
586        if isinstance(response_wire, str):
587            map_func = lambda x: ord(x)
588        else:
589            map_func = lambda x: x
590
591        timeouts = self._get_num_timeouts(response)
592        is_timeout = isinstance(response, dns.exception.Timeout)
593        is_valid = isinstance(response, dns.message.Message) and response.rcode() in (dns.rcode.NOERROR, dns.rcode.NXDOMAIN)
594        is_truncated = response_wire is not None and map_func(response_wire[2]) & 0x02
595        if response_wire is not None:
596            response_len = len(response_wire)
597        else:
598            response_len = None
599
600        if self._request.edns >= 0 and \
601                (is_timeout or is_valid or is_truncated):
602            pass
603        else:
604            self._state = self.INVALID
605            return
606
607        if self._state == self.START:
608            if timeouts >= self._initial_timeouts:
609                self._lower_bound = self._reduced_payload
610                self._upper_bound = self._request.payload - 1
611                self._request.payload = self._reduced_payload
612                self._state = self.REDUCED_PAYLOAD
613                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, self._reduced_payload)
614
615        elif self._state == self.REDUCED_PAYLOAD:
616            if timeouts >= self._max_timeouts:
617                self._state == self.INVALID
618                return None
619
620            if not is_timeout:
621                if is_truncated or is_valid:
622                    self._lower_bound = self._water_mark = response_len
623                    self._params['timeout'] = self._bounding_timeout
624                    self._params['tcp'] = True
625                    self._state = self.USE_TCP
626                    if is_truncated:
627                        return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TC_SET, response_len, RETRY_ACTION_USE_TCP, None)
628                    else:
629                        return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, response_len, RETRY_ACTION_USE_TCP, None)
630
631        elif self._state == self.USE_TCP:
632            if not is_timeout and is_valid:
633                #XXX this is cheating because we're not reporting the change to UDP
634                self._params['tcp'] = False
635                payload = response_len - 1
636                self._request.payload = payload
637                self._state = self.TCP_MINUS_ONE
638                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, response_len, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
639
640        elif self._state == self.TCP_MINUS_ONE:
641            if is_timeout:
642                self._upper_bound = self._request.payload - 1
643                payload = self._lower_bound + (self._upper_bound + 1 - self._lower_bound)//2
644                self._request.payload = payload
645                self._state = self.PICKLE
646                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
647            # if the size of the message is less than the watermark, then perhaps we were rate limited
648            elif response_wire is not None and response_len < self._water_mark:
649                # but if this isn't the first time, just quit.  it could be that
650                # the server simply has some wonky way of determining how/where to truncate.
651                if self._history[-1].cause == RETRY_CAUSE_DIAGNOSTIC and self._history[-1].action == RETRY_ACTION_CHANGE_SPORT:
652                    self._params['tcp'] = True
653                    self._state = self.TCP_FINAL
654                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, None, RETRY_ACTION_USE_TCP, None)
655                else:
656                    self._params['wait'] = 1.0
657                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, None, RETRY_ACTION_CHANGE_SPORT, None)
658            # if the response was truncated, then the size of the payload
659            # received via TCP is the largest we can receive
660            elif is_truncated:
661                self._params['tcp'] = True
662                self._state = self.TCP_FINAL
663                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TC_SET, response_len, RETRY_ACTION_USE_TCP, None)
664
665        elif self._state == self.PICKLE:
666            if self._upper_bound - self._lower_bound <= 1:
667                self._params['tcp'] = True
668                self._state = self.TCP_FINAL
669                if is_truncated:
670                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TC_SET, response_len, RETRY_ACTION_USE_TCP, None)
671                elif is_timeout:
672                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_USE_TCP, None)
673                elif not is_valid:
674                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, None, RETRY_ACTION_USE_TCP, None)
675            elif is_timeout:
676                self._upper_bound = self._request.payload - 1
677                payload = self._lower_bound + (self._upper_bound + 1 - self._lower_bound)//2
678                self._request.payload = payload
679                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
680            # if the size of the message is less than the watermark, then perhaps we were rate limited
681            elif response_len < self._water_mark:
682                # but if this isn't the first time, just quit.  it could be that
683                # the server simply has some wonky way of determining how/where to truncate.
684                if self._history[-1].cause == RETRY_CAUSE_DIAGNOSTIC and self._history[-1].action == RETRY_ACTION_CHANGE_SPORT:
685                    self._params['tcp'] = True
686                    self._state = self.TCP_FINAL
687                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, None, RETRY_ACTION_USE_TCP, None)
688                else:
689                    self._params['wait'] = 1.0
690                    return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, None, RETRY_ACTION_CHANGE_SPORT, None)
691            elif is_valid:
692                self._lower_bound = self._request.payload
693                payload = self._lower_bound + (self._upper_bound + 1 - self._lower_bound)//2
694                self._request.payload = payload
695                return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, response_len, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
696
697        elif self._state == self.TCP_FINAL:
698            pass
699
700        elif self._state == self.INVALID:
701            pass
702
703class ChangeTimeoutOnTimeoutHandler(ActionIndependentDNSResponseHandler):
704    '''Modify timeout value when a certain number of timeouts is reached.'''
705
706    def __init__(self, timeout, timeouts):
707        self._timeout = timeout
708        self._timeouts = timeouts
709
710    def handle(self, response_wire, response, response_time):
711        timeouts = self._get_num_timeouts(response)
712        if isinstance(response, dns.exception.Timeout) and timeouts == self._timeouts:
713            self._params['timeout'] = self._timeout
714
715class RetryOnTimeoutHandler(DNSResponseHandler):
716    '''Retry with no change when a query times out.'''
717
718    def handle(self, response_wire, response, response_time):
719        if isinstance(response, dns.exception.Timeout):
720            if self._params['tcp']:
721                action = RETRY_ACTION_CHANGE_SPORT
722            else:
723                action = RETRY_ACTION_NO_CHANGE
724            return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, action, None)
725
726class DefaultAcceptHandler(DNSResponseHandler):
727    '''Accept the response if there was no other reason to not accept it.'''
728
729    def handle(self, response_wire, response, response_time):
730        raise AcceptResponse()
731
732class LifetimeHandler(ActionIndependentDNSResponseHandler):
733    '''Stop handling and retrying if the designated lifetime has been
734    exceeded.'''
735
736    def __init__(self, lifetime):
737        self._lifetime = lifetime
738        self._start = time.time()
739
740    def handle(self, response_wire, response, response_time):
741        if self.time_remaining() <= 0:
742            raise AcceptResponse()
743
744    def time_remaining(self):
745        return max(self._start + self._lifetime - time.time(), 0)
746
747class MaxTimeoutsHandler(ActionIndependentDNSResponseHandler):
748    '''Stop handling and retrying if the maximum number of timeouts has been
749    exceeded.'''
750
751    def __init__(self, max_timeouts):
752        self._max_timeouts = max_timeouts
753
754    def handle(self, response_wire, response, response_time):
755        if self._get_num_timeouts(response) >= self._max_timeouts:
756            raise AcceptResponse()
757
758class DNSQueryHandler:
759    '''A handler associated with a DNS query to a server.'''
760
761    def __init__(self, query, request, server_cookie, server_cookie_status, params, response_handlers, server, client):
762        self.query = query
763        self.request = request
764        self.params = params
765        self.server_cookie = server_cookie
766        self.server_cookie_status = server_cookie_status
767        self._response_handlers = response_handlers
768        self.history = []
769        self._server = server
770        self._client = client
771
772        for handler in self._response_handlers:
773            handler.set_context(self.params, self.history, self.request)
774
775        if query.lifetime is not None:
776            self._expiration = time.time() + query.lifetime
777        else:
778            self._expiration = None
779
780        self._set_query_time()
781
782    def _set_query_time(self):
783        self.query_time = time.time() + self.params['wait']
784
785    def _reset_wait(self):
786        self.params['wait'] = 0
787
788    def get_query_transport_meta(self):
789        return transport.DNSQueryTransportMeta(self.request.to_wire(), self._server, self.params['tcp'], self.get_timeout(), \
790                self.query.odd_ports.get(self._server, self.query.port), src=self._client, sport=self.params['sport'])
791
792    def get_remaining_lifetime(self):
793        if self._expiration is None:
794            # send arbitrarily high value
795            return 86400
796        return max(self._expiration - time.time(), 0)
797
798    def get_timeout(self):
799        if self._expiration is None:
800            return self.params['timeout']
801        timeout = min(self.params['timeout'], self.get_remaining_lifetime())
802        if timeout < MIN_QUERY_TIMEOUT:
803            return MIN_QUERY_TIMEOUT
804        return timeout
805
806    def handle_response(self, response_wire, response, response_time, client, sport):
807        retry_action = None
808        try:
809            for handler in self._response_handlers:
810                if retry_action is None:
811                    retry_action = handler.handle(response_wire, response, response_time)
812                    if retry_action is not None:
813                        if retry_action.action == RETRY_ACTION_NO_CHANGE:
814                            self.params['sport'] = sport
815                        else:
816                            self.params['sport'] = None
817                elif isinstance(handler, ActionIndependentDNSResponseHandler):
818                    handler.handle(response_wire, response, response_time)
819
820            if retry_action is not None:
821                # If we were unable to bind to the source address, then this is
822                # our fault
823                if retry_action.cause == RETRY_CAUSE_NETWORK_ERROR and retry_action.cause_arg == errno.EADDRNOTAVAIL:
824                    raise AcceptResponse
825
826                # If there is no client-side connectivity, then simply return.
827                #
828                #XXX (Note that this only catches the case when a client IP has
829                # not been explicitly specified (i.e., self._client is None).
830                # Explicitly specifying a client IP that cannot connect to a
831                # given destination (e.g., because it is of the wrong address
832                # scope) will result in a regular network failure with
833                # EHOSTUNREACH or ENETUNREACH, as there is no scope comparison
834                # in this code.)
835                if retry_action.cause == RETRY_CAUSE_NETWORK_ERROR and retry_action.cause_arg in (errno.EHOSTUNREACH, errno.ENETUNREACH, errno.EAFNOSUPPORT) and client is None:
836                    raise AcceptResponse
837
838                # if this error was our fault, don't add it to the history
839                if retry_action.cause == RETRY_CAUSE_NETWORK_ERROR and retry_action.cause_arg == errno.EMFILE:
840                    pass
841                else:
842                    self.history.append(retry_action)
843
844            self._set_query_time()
845            self._reset_wait()
846
847        except AcceptResponse:
848            return response
849
850class AggregateDNSResponse(object):
851    ttl_cmp = False
852
853    def __init__(self):
854        self.answer_info = []
855        self.nodata_info = []
856        self.nxdomain_info = []
857        self.referral_info = []
858        self.truncated_info = []
859        self.error_info = []
860
861    def _aggregate_response(self, server, client, response, qname, rdtype, rdclass, bailiwick):
862        if response.is_valid_response():
863            if response.is_complete_response():
864                is_referral = response.is_referral(qname, rdtype, rdclass, bailiwick)
865                self._aggregate_answer(server, client, response, is_referral, qname, rdtype, rdclass)
866            else:
867                truncated_info = TruncatedResponse(response.message.to_wire())
868                DNSResponseComponent.insert_into_list(truncated_info, self.truncated_info, server, client, response)
869
870        else:
871            self._aggregate_error(server, client, response)
872
873    def _aggregate_answer(self, server, client, response, referral, qname, rdtype, rdclass):
874        msg = response.message
875
876        # sort with the most specific DNAME infos first
877        dname_rrsets = [x for x in msg.answer if x.rdtype == dns.rdatatype.DNAME and x.rdclass == rdclass]
878        dname_rrsets.sort(reverse=True)
879
880        qname_sought = qname
881        try:
882            i = 0
883            while i < MAX_CNAME_REDIRECTION:
884
885                # synthesize a CNAME from a DNAME, if possible
886                synthesized_cname_info = None
887                for dname_rrset in dname_rrsets:
888                    if qname_sought.parent().is_subdomain(dname_rrset.name):
889                        synthesized_cname_info = RRsetInfo(cname_from_dname(qname_sought, dname_rrset), self.ttl_cmp, RRsetInfo(dname_rrset, self.ttl_cmp))
890                        break
891
892                try:
893                    rrset_info = self._aggregate_answer_rrset(server, client, response, qname_sought, rdtype, rdclass, referral)
894
895                    # if there was a synthesized CNAME, add it to the rrset_info
896                    if rrset_info.rrset.rdtype == dns.rdatatype.CNAME and rrset_info.rrset.rdclass == rdclass and synthesized_cname_info is not None:
897                        synthesized_cname_info = rrset_info.create_or_update_cname_from_dname_info(synthesized_cname_info, server, client, response, rdclass)
898                        synthesized_cname_info.update_rrsig_info(server, client, response, msg.answer, rdclass, referral)
899
900                except KeyError:
901                    if synthesized_cname_info is None:
902                        raise
903                    synthesized_cname_info = DNSResponseComponent.insert_into_list(synthesized_cname_info, self.answer_info, server, client, response)
904                    synthesized_cname_info.dname_info.update_rrsig_info(server, client, response, msg.answer, rdclass, referral)
905                    rrset_info = synthesized_cname_info
906
907                if rrset_info.rrset.rdtype == dns.rdatatype.CNAME and rrset_info.rrset.rdclass == rdclass:
908                    qname_sought = rrset_info.rrset[0].target
909                else:
910                    break
911                i += 1
912        except KeyError:
913            if referral and rdtype != dns.rdatatype.DS:
914                # add referrals
915                try:
916                    rrset = [x for x in msg.authority if qname.is_subdomain(x.name) and x.rdtype == dns.rdatatype.NS and x.rdclass == rdclass][0]
917                except IndexError:
918                    pass
919                else:
920                    referral_info = ReferralResponse(rrset.name)
921                    DNSResponseComponent.insert_into_list(referral_info, self.referral_info, server, client, response)
922
923                # with referrals, don't do any further processing
924                return
925
926            # don't store no answer or NXDOMAIN info for names other than qname
927            # if recursion is not desired and available
928            if qname_sought != qname and not response.recursion_desired_and_available():
929                return
930
931            if msg.rcode() == dns.rcode.NXDOMAIN:
932                neg_response_info_list = self.nxdomain_info
933            else:
934                neg_response_info_list = self.nodata_info
935
936            neg_response_info = NegativeResponseInfo(qname_sought, rdtype, self.ttl_cmp)
937            neg_response_info = DNSResponseComponent.insert_into_list(neg_response_info, neg_response_info_list, server, client, response)
938            neg_response_info.create_or_update_nsec_info(server, client, response, rdclass, referral)
939            neg_response_info.create_or_update_soa_info(server, client, response, rdclass, referral)
940
941    def _aggregate_answer_rrset(self, server, client, response, qname, rdtype, rdclass, referral):
942        msg = response.message
943
944        try:
945            rrset = msg.find_rrset(msg.answer, qname, rdclass, rdtype)
946        except KeyError:
947            rrset = msg.find_rrset(msg.answer, qname, rdclass, dns.rdatatype.CNAME)
948
949        rrset_info = RRsetInfo(rrset, self.ttl_cmp)
950        rrset_info = DNSResponseComponent.insert_into_list(rrset_info, self.answer_info, server, client, response)
951
952        rrset_info.update_rrsig_info(server, client, response, msg.answer, rdclass, referral)
953
954        return rrset_info
955
956    def _aggregate_error(self, server, client, response):
957        msg = response.message
958        if msg is None:
959            error_info = DNSResponseError(response.error, response.errno)
960        else:
961            error_info = DNSResponseError(RESPONSE_ERROR_INVALID_RCODE, msg.rcode())
962        error_info = DNSResponseComponent.insert_into_list(error_info, self.error_info, server, client, response)
963
964class DNSQuery(object):
965    '''An simple DNS Query and its responses.'''
966
967    def __init__(self, qname, rdtype, rdclass,
968            flags, edns, edns_max_udp_payload, edns_flags, edns_options, tcp):
969
970        self.qname = qname
971        self.rdtype = rdtype
972        self.rdclass = rdclass
973        self.flags = flags
974        self.edns = edns
975        self.edns_max_udp_payload = edns_max_udp_payload
976        self.edns_flags = edns_flags
977        self.edns_options = edns_options
978        self.tcp = tcp
979
980        self.responses = {}
981
982    def copy(self, bailiwick_map, default_bailiwick, with_responses=True):
983        '''Return a clone of the current DNSQuery instance.  Parameters are
984        passed by reference rather than copied.  Note: if it turns out that
985        these member variables might be modified somehow by other instances in
986        future use, then these will need to be copies.'''
987
988        clone = DNSQuery(self.qname, self.rdtype, self.rdclass,
989                self.flags, self.edns, self.edns_max_udp_payload, self.edns_flags, self.edns_options, self.tcp)
990
991        if with_responses:
992            for server in self.responses:
993                bailiwick = bailiwick_map.get(server, default_bailiwick)
994                for client, response in self.responses[server].items():
995                    response_clone = response.copy()
996                    response_clone.query = clone
997                    clone.add_response(server, client, response_clone, bailiwick)
998
999        return clone
1000
1001    def join(self, query, bailiwick_map, default_bailiwick):
1002        if not (isinstance(query, DNSQuery)):
1003            raise ValueError('A DNSQuery instance can only be joined with another DNSQuery instance.')
1004
1005        if not (self.qname.to_text() == query.qname.to_text() and self.rdtype == query.rdtype and \
1006                self.rdclass == query.rdclass and self.flags == query.flags and \
1007                self.edns == query.edns and self.edns_max_udp_payload == query.edns_max_udp_payload and \
1008                self.edns_flags == query.edns_flags and self.edns_options == query.edns_options and \
1009                self.tcp == query.tcp):
1010            raise ValueError('DNS query parameters for DNSQuery instances being joined must be the same.')
1011
1012        clone = self.copy(bailiwick_map, default_bailiwick)
1013        for server in query.responses:
1014            bailiwick = bailiwick_map.get(server, default_bailiwick)
1015            for client, response in query.responses[server].items():
1016                response_clone = response.copy()
1017                response_clone.query = clone
1018                clone.add_response(server, client, response_clone, bailiwick)
1019        return clone
1020
1021    def project(self, servers, bailiwick_map, default_bailiwick):
1022        if servers.difference(self.responses):
1023            raise ValueError('A DNSQuery can only project responses from servers that have been queried.')
1024
1025        clone = self.copy(bailiwick_map, default_bailiwick, with_responses=False)
1026        for server in servers:
1027            bailiwick = bailiwick_map.get(server, default_bailiwick)
1028            for client, response in self.responses[server].items():
1029                response_clone = response.copy()
1030                response_clone.query = clone
1031                clone.add_response(server, client, response_clone, bailiwick)
1032        return clone
1033
1034    def add_response(self, server, client, response, bailiwick):
1035        if server not in self.responses:
1036            self.responses[server] = {}
1037        if response.query is not None and response.query is not self:
1038            raise ValueError('Response for %s/%s is already associated with a query.' % (self.qname, dns.rdatatype.to_text(self.rdtype)))
1039        if client in self.responses[server]:
1040            raise ValueError('Response for %s/%s from server %s to client %s already exists.' % (self.qname, dns.rdatatype.to_text(self.rdtype), server, client))
1041        response.query = self
1042        self.responses[server][client] = response
1043
1044    def is_authoritative_answer_all(self):
1045        val = None
1046        for server in self.responses:
1047            for response in self.responses[server].values():
1048                if not (response.is_valid_response() and response.is_complete_response()):
1049                    continue
1050                if response.is_authoritative() and response.is_answer(self.qname, self.rdtype):
1051                    val = True
1052                else:
1053                    return False
1054
1055        if val is None:
1056            val = False
1057        return val
1058
1059    def is_answer_any(self):
1060        for server in self.responses:
1061            for response in self.responses[server].values():
1062                if not (response.is_valid_response() and response.is_complete_response()):
1063                    continue
1064                if response.is_answer(self.qname, self.rdtype):
1065                    return True
1066        return False
1067
1068    def is_nxdomain_all(self):
1069        val = None
1070        for server in self.responses:
1071            for response in self.responses[server].values():
1072                if not (response.is_valid_response() and response.is_complete_response()):
1073                    continue
1074                if response.is_nxdomain(self.qname, self.rdtype):
1075                    if val is None:
1076                        val = True
1077                else:
1078                    return False
1079
1080        if val is None:
1081            val = False
1082        return val
1083
1084    def is_not_delegation_all(self):
1085        val = None
1086        for server in self.responses:
1087            for response in self.responses[server].values():
1088                if not (response.is_valid_response() and response.is_complete_response()):
1089                    continue
1090                if response.not_delegation(self.qname, self.rdtype):
1091                    if val is None:
1092                        val = True
1093                else:
1094                    return False
1095
1096        if val is None:
1097            val = False
1098        return val
1099
1100    def is_valid_complete_response_any(self):
1101        for server in self.responses:
1102            for response in self.responses[server].values():
1103                if response.is_valid_response() and response.is_complete_response():
1104                    return True
1105        return False
1106
1107    def is_valid_complete_authoritative_response_any(self):
1108        for server in self.responses:
1109            for response in self.responses[server].values():
1110                if response.is_valid_response() and response.is_complete_response() and response.is_authoritative():
1111                    return True
1112        return False
1113
1114    def servers_with_valid_complete_response(self, bailiwick_map, default_bailiwick):
1115        servers_clients = set()
1116        for server in self.responses:
1117            bailiwick = bailiwick_map.get(server, default_bailiwick)
1118            for client, response in self.responses[server].items():
1119                if response.is_valid_response() and response.is_complete_response() and not response.is_referral(self.qname, self.rdtype, self.rdclass, bailiwick):
1120                    servers_clients.add((server, client))
1121        return servers_clients
1122
1123    def is_nxdomain_any(self):
1124        for server in self.responses:
1125            for response in self.responses[server].values():
1126                if not (response.is_valid_response() and response.is_complete_response()):
1127                    continue
1128                if response.is_nxdomain(self.qname, self.rdtype):
1129                    return True
1130        return False
1131
1132    def serialize(self, meta_only=False):
1133        d = OrderedDict((
1134            ('qname', lb2s(self.qname.to_text())),
1135            ('qclass', dns.rdataclass.to_text(self.rdclass)),
1136            ('qtype', dns.rdatatype.to_text(self.rdtype)),
1137        ))
1138        d['options'] = OrderedDict((
1139            ('flags', self.flags),
1140        ))
1141        if self.edns >= 0:
1142            d['options']['edns_version'] = self.edns
1143            d['options']['edns_max_udp_payload'] = self.edns_max_udp_payload
1144            d['options']['edns_flags'] = self.edns_flags
1145            d['options']['edns_options'] = []
1146            for o in self.edns_options:
1147                s = io.BytesIO()
1148                o.to_wire(s)
1149                d['options']['edns_options'].append((o.otype, lb2s(binascii.hexlify(s.getvalue()))))
1150            d['options']['tcp'] = self.tcp
1151
1152        d['responses'] = OrderedDict()
1153        servers = list(self.responses.keys())
1154        servers.sort()
1155        for server in servers:
1156            d['responses'][server] = OrderedDict()
1157            clients = list(self.responses[server].keys())
1158            clients.sort()
1159            for client in clients:
1160                if meta_only:
1161                    d['responses'][server][client] = self.responses[server][client].serialize_meta()
1162                else:
1163                    d['responses'][server][client] = self.responses[server][client].serialize()
1164
1165        return d
1166
1167    @classmethod
1168    def deserialize(self, d, bailiwick_map, default_bailiwick, cookie_jar_map, default_cookie_jar, cookie_standin, cookie_bad):
1169        qname = dns.name.from_text(d['qname'])
1170        rdclass = dns.rdataclass.from_text(d['qclass'])
1171        rdtype = dns.rdatatype.from_text(d['qtype'])
1172
1173        d1 = d['options']
1174
1175        flags = d1['flags']
1176        if 'edns_version' in d1:
1177            edns = d1['edns_version']
1178            edns_max_udp_payload = d1['edns_max_udp_payload']
1179            edns_flags = d1['edns_flags']
1180            edns_options = []
1181            for otype, data in d1['edns_options']:
1182                edns_options.append(dns.edns.GenericOption(otype, binascii.unhexlify(data)))
1183        else:
1184            edns = None
1185            edns_max_udp_payload = None
1186            edns_flags = None
1187            edns_options = []
1188
1189        tcp = d1['tcp']
1190
1191        q = DNSQuery(qname, rdtype, rdclass,
1192                flags, edns, edns_max_udp_payload, edns_flags, edns_options, tcp)
1193
1194        server_cookie = None
1195        server_cookie_status = DNS_COOKIE_NO_COOKIE
1196        if edns >= 0:
1197            try:
1198                cookie_opt = [o for o in edns_options if o.otype == 10][0]
1199            except IndexError:
1200                pass
1201            else:
1202                if len(cookie_opt.data) == 8:
1203                    server_cookie_status = DNS_COOKIE_CLIENT_COOKIE_ONLY
1204                elif len(cookie_opt.data) >= 16 and len(cookie_opt.data) <= 40:
1205                    if cookie_opt.data[8:] == cookie_standin:
1206                        # initially assume that there is a cookie for the server;
1207                        # change the value later if there isn't
1208                        server_cookie_status = DNS_COOKIE_SERVER_COOKIE_FRESH
1209                    elif cookie_opt.data[8:] == cookie_bad:
1210                        server_cookie_status = DNS_COOKIE_SERVER_COOKIE_BAD
1211                    else:
1212                        server_cookie_status = DNS_COOKIE_SERVER_COOKIE_STATIC
1213                else:
1214                    server_cookie_status = DNS_COOKIE_IMPROPER_LENGTH
1215
1216        for server in d['responses']:
1217            server_ip = IPAddr(server)
1218            bailiwick = bailiwick_map.get(server_ip, default_bailiwick)
1219            cookie_jar = cookie_jar_map.get(server_ip, default_cookie_jar)
1220            server_cookie = cookie_jar.get(server_ip, None)
1221            status = server_cookie_status
1222            if status == DNS_COOKIE_SERVER_COOKIE_FRESH and server_cookie is None:
1223                status = DNS_COOKIE_CLIENT_COOKIE_ONLY
1224            for client in d['responses'][server]:
1225                q.add_response(server_ip, IPAddr(client), DNSResponse.deserialize(d['responses'][server][client], q, server_cookie, status), bailiwick)
1226        return q
1227
1228class DNSQueryAggregateDNSResponse(DNSQuery, AggregateDNSResponse):
1229    def __init__(self, qname, rdtype, rdclass,
1230            flags, edns, edns_max_udp_payload, edns_flags, edns_options, tcp):
1231        DNSQuery.__init__(self, qname, rdtype, rdclass,
1232            flags, edns, edns_max_udp_payload, edns_flags, edns_options, tcp)
1233        AggregateDNSResponse.__init__(self)
1234
1235    def add_response(self, server, client, response, bailiwick):
1236        super(DNSQueryAggregateDNSResponse, self).add_response(server, client, response, bailiwick)
1237        self._aggregate_response(server, client, response, self.qname, self.rdtype, self.rdclass, bailiwick)
1238
1239class MultiQuery(object):
1240    '''An simple DNS Query and its responses.'''
1241
1242    def __init__(self, qname, rdtype, rdclass):
1243        self.qname = qname
1244        self.rdtype = rdtype
1245        self.rdclass = rdclass
1246
1247        self.queries = {}
1248
1249    def add_query(self, query, bailiwick_map, default_bailiwick):
1250        if not (self.qname == query.qname and self.rdtype == query.rdtype and self.rdclass == query.rdclass):
1251            raise ValueError('DNS query information must be the same as that to which query is being joined.')
1252
1253        edns_options_str = b''
1254        for o in query.edns_options:
1255            s = io.BytesIO()
1256            o.to_wire(s)
1257            edns_options_str += struct.pack(b'!H', o.otype) + s.getvalue()
1258        params = (query.qname.to_text(), query.flags, query.edns, query.edns_max_udp_payload, query.edns_flags, edns_options_str, query.tcp)
1259        if params in self.queries:
1260            self.queries[params] = self.queries[params].join(query, bailiwick_map, default_bailiwick)
1261        else:
1262            self.queries[params] = query
1263
1264    def project(self, servers, bailiwick_map, default_bailiwick):
1265        query = self.__class__(self.qname, self.rdtype, self.rdclass)
1266
1267        for params in self.queries:
1268            query.add_query(self.queries[params].project(servers, bailiwick_map, default_bailiwick))
1269        return query
1270
1271    def is_nxdomain_all(self):
1272        for params in self.queries:
1273            if not self.queries[params].is_nxdomain_all():
1274                return False
1275        return True
1276
1277    def is_valid_complete_authoritative_response_any(self):
1278        for params in self.queries:
1279            if self.queries[params].is_valid_complete_authoritative_response_any():
1280                return True
1281        return False
1282
1283class MultiQueryAggregateDNSResponse(MultiQuery, AggregateDNSResponse):
1284    def __init__(self, qname, rdtype, rdclass):
1285        MultiQuery.__init__(self, qname, rdtype, rdclass)
1286        AggregateDNSResponse.__init__(self)
1287
1288    def add_query(self, query, bailiwick_map, default_bailiwick):
1289        super(MultiQueryAggregateDNSResponse, self).add_query(query, bailiwick_map, default_bailiwick)
1290        for server in query.responses:
1291            bailiwick = bailiwick_map.get(server, default_bailiwick)
1292            for client, response in query.responses[server].items():
1293                self._aggregate_response(server, client, response, self.qname, self.rdtype, self.rdclass, bailiwick)
1294
1295class TTLDistinguishingMultiQueryAggregateDNSResponse(MultiQueryAggregateDNSResponse):
1296    ttl_cmp = True
1297
1298class ExecutableDNSQuery(DNSQuery):
1299    '''An executable DNS Query.'''
1300
1301    default_th_factory = transport.DNSQueryTransportHandlerDNSPrivateFactory()
1302
1303    def __init__(self, qname, rdtype, rdclass, servers, bailiwick,
1304            client_ipv4, client_ipv6, port, odd_ports, cookie_jar, cookie_standin, cookie_bad,
1305            flags, edns, edns_max_udp_payload, edns_flags, edns_options, tcp,
1306            response_handlers, query_timeout, max_attempts, lifetime):
1307
1308        super(ExecutableDNSQuery, self).__init__(qname, rdtype, rdclass,
1309                flags, edns, edns_max_udp_payload, edns_flags, edns_options, tcp)
1310
1311        if not isinstance(servers, set):
1312            if isinstance(servers, (list, tuple)):
1313                servers = set(servers)
1314            else:
1315                servers = set([servers])
1316        if not servers:
1317            raise ValueError("At least one server must be specified for an ExecutableDNSQuery")
1318
1319        self.servers = servers
1320        self.bailiwick = bailiwick
1321        self.client_ipv4 = client_ipv4
1322        self.client_ipv6 = client_ipv6
1323        self.port = port
1324        if odd_ports is None:
1325            odd_ports = {}
1326        self.odd_ports = odd_ports
1327        if cookie_jar is None:
1328            cookie_jar = {}
1329        self.cookie_jar = cookie_jar
1330        self.cookie_standin = cookie_standin
1331        self.cookie_bad = cookie_bad
1332        self.response_handlers = response_handlers
1333
1334        self.query_timeout = query_timeout
1335
1336        if lifetime is None and max_attempts is None:
1337            raise ValueError("At least one of lifetime or max_attempts must be specified for an ExecutableDNSQuery instance.")
1338        self.max_attempts = max_attempts
1339        self.lifetime = lifetime
1340
1341        self._executed = False
1342
1343    def get_query_handler(self, server):
1344        edns_options = copy.deepcopy(self.edns_options)
1345        server_cookie = None
1346        server_cookie_status = DNS_COOKIE_NO_COOKIE
1347
1348        if self.edns >= 0:
1349            try:
1350                cookie_opt = [o for o in edns_options if o.otype == 10][0]
1351            except IndexError:
1352                pass
1353            else:
1354                if len(cookie_opt.data) == 8:
1355                    server_cookie_status = DNS_COOKIE_CLIENT_COOKIE_ONLY
1356                elif len(cookie_opt.data) >= 16 and len(cookie_opt.data) <= 40:
1357                    if cookie_opt.data[8:] == self.cookie_standin:
1358                        if server in self.cookie_jar:
1359                            # if there is a cookie for this server,
1360                            # then add it
1361                            server_cookie = self.cookie_jar[server]
1362                            cookie_opt.data = cookie_opt.data[:8] + server_cookie
1363                            server_cookie_status = DNS_COOKIE_SERVER_COOKIE_FRESH
1364                        else:
1365                            # otherwise, send just the client cookie.
1366                            cookie_opt.data = cookie_opt.data[:8]
1367                            server_cookie_status = DNS_COOKIE_CLIENT_COOKIE_ONLY
1368                    elif cookie_opt.data[8:] == self.cookie_bad:
1369                        server_cookie_status = DNS_COOKIE_SERVER_COOKIE_BAD
1370                    else:
1371                        server_cookie_status = DNS_COOKIE_SERVER_COOKIE_STATIC
1372                else:
1373                    server_cookie_status = DNS_COOKIE_IMPROPER_LENGTH
1374
1375        request = dns.message.Message()
1376        request.flags = self.flags
1377        request.find_rrset(request.question, self.qname, self.rdclass, self.rdtype, create=True, force_unique=True)
1378        request.use_edns(self.edns, self.edns_flags, self.edns_max_udp_payload, options=edns_options)
1379
1380        if server.version == 6:
1381            client = self.client_ipv6
1382        else:
1383            client = self.client_ipv4
1384
1385        params = { 'tcp': self.tcp, 'sport': None, 'wait': 0, 'timeout': self.query_timeout }
1386
1387        response_handlers = [RetryOnNetworkErrorHandler(3).build()] + [h.build() for h in self.response_handlers] + \
1388            [RetryOnTimeoutHandler().build(), DefaultAcceptHandler().build()]
1389
1390        if self.max_attempts is not None:
1391            response_handlers.append(MaxTimeoutsHandler(self.max_attempts).build())
1392        if self.lifetime is not None:
1393            response_handlers.append(LifetimeHandler(self.lifetime).build())
1394
1395        return DNSQueryHandler(self, request, server_cookie, server_cookie_status, params, response_handlers, server, client)
1396
1397    @classmethod
1398    def execute_queries(cls, *queries, **kwargs):
1399        '''Execute the query to a given server, and handle it appropriately.'''
1400
1401        tm = kwargs.get('tm', None)
1402        if tm is None:
1403            # this starts a thread that stops when tm goes out of scope
1404            tm = transport.DNSQueryTransportManager()
1405
1406        th_factories = kwargs.get('th_factories', None)
1407        if th_factories is None:
1408            th_factories = (cls.default_th_factory,)
1409
1410        request_list = []
1411        response_queue = queue.Queue()
1412
1413        ignore_queryid = kwargs.get('ignore_queryid', True)
1414        response_wire_map = {}
1415
1416        query_handlers = {}
1417        query_time = None
1418        for th_factory in th_factories:
1419            if not th_factory.cls.singleton:
1420                th = th_factory.build(processed_queue=response_queue)
1421
1422            for query in queries:
1423                qtm_for_server = False
1424                for server in query.servers:
1425                    if not th_factory.cls.allow_loopback_query and (LOOPBACK_IPV4_RE.match(server) or server == LOOPBACK_IPV6):
1426                        continue
1427                    if not th_factory.cls.allow_private_query and (RFC_1918_RE.match(server) or LINK_LOCAL_RE.match(server) or UNIQ_LOCAL_RE.match(server)):
1428                        continue
1429
1430                    qtm_for_server = True
1431                    qh = query.get_query_handler(server)
1432                    qtm = qh.get_query_transport_meta()
1433                    query_handlers[qtm] = qh
1434
1435                    if th_factory.cls.singleton:
1436                        th = th_factory.build(processed_queue=response_queue)
1437                        th.add_qtm(qtm)
1438                        th.init_req()
1439                        bisect.insort(request_list, (qh.query_time, th))
1440                    else:
1441                        # find the maximum query time
1442                        if query_time is None or qh.query_time > query_time:
1443                            query_time = qh.query_time
1444                        th.add_qtm(qtm)
1445
1446                if not qtm_for_server:
1447                    raise NoValidServersToQuery('No valid servers to query!')
1448
1449            if not th_factory.cls.singleton:
1450                th.init_req()
1451                bisect.insort(request_list, (query_time, th))
1452
1453        while query_handlers:
1454            while request_list and time.time() >= request_list[0][0]:
1455                tm.handle_msg_nowait(request_list.pop(0)[1])
1456
1457            t = time.time()
1458            if request_list and t < request_list[0][0]:
1459                timeout = max(request_list[0][0] - t, 0)
1460            else:
1461                timeout = None
1462
1463            try:
1464                # pull a response from the queue
1465                th = response_queue.get(timeout=timeout)
1466            except queue.Empty:
1467                continue
1468            th.finalize()
1469
1470            newth = th.factory.build(processed_queue=response_queue)
1471            query_time = None
1472            for qtm in th.qtms:
1473                # find its matching query meta information
1474                qh = query_handlers.pop(qtm)
1475                query = qh.query
1476
1477                # define response as either a Message created from parsing
1478                # the wire response or an Exception
1479                if qtm.err is not None:
1480                    response = qtm.err
1481                else:
1482                    wire_zero_queryid = b'\x00\x00' + qtm.res[2:]
1483                    if wire_zero_queryid in response_wire_map:
1484                        response = response_wire_map[wire_zero_queryid]
1485                    else:
1486                        try:
1487                            response = dns.message.from_wire(qtm.res)
1488                        except Exception as e:
1489                            response = e
1490                        if ignore_queryid:
1491                            response_wire_map[wire_zero_queryid] = response
1492                if qtm.res:
1493                    msg_size = len(qtm.res)
1494                else:
1495                    msg_size = None
1496                response_time = round(qtm.end_time - qtm.start_time, 3)
1497                response = qh.handle_response(qtm.res, response, response_time, qtm.src, qtm.sport)
1498
1499                # if no response was returned, then resubmit the modified query
1500                if response is None:
1501                    qtm = qh.get_query_transport_meta()
1502                    # find the maximum query time
1503                    if query_time is None or qh.query_time > query_time:
1504                        query_time = qh.query_time
1505                    query_handlers[qtm] = qh
1506                    newth.add_qtm(qtm)
1507                    continue
1508
1509                # otherwise store away the response (or error), history, and response time
1510                if isinstance(response, dns.message.Message):
1511                    msg = response
1512                    err = None
1513                    errno1 = None
1514                else:
1515                    msg = None
1516                    if isinstance(response, dns.exception.Timeout):
1517                        err = RESPONSE_ERROR_TIMEOUT
1518                    elif isinstance(response, (socket.error, EOFError)):
1519                        err = RESPONSE_ERROR_NETWORK_ERROR
1520                    elif isinstance(response, (struct.error, dns.exception.FormError)):
1521                        err = RESPONSE_ERROR_FORMERR
1522                    #XXX need to determine how to handle non-parsing
1523                    # validation errors with dnspython (e.g., signature with
1524                    # no keyring)
1525                    else:
1526                        err = RESPONSE_ERROR_OTHER
1527                    if hasattr(response, 'errno'):
1528                        errno1 = response.errno
1529                    else:
1530                        errno1 = None
1531                response_obj = DNSResponse(msg, msg_size, err, errno1, qh.history, response_time, query, qh.server_cookie, qh.server_cookie_status)
1532
1533                # if client IP is not specified, and there is a socket
1534                # failure, then src might be None
1535                if qtm.src is not None:
1536                    src = IPAddr(qtm.src)
1537                else:
1538                    src = qtm.src
1539
1540                # If this was a network error, determine if it was a binding
1541                # error
1542                if err == RESPONSE_ERROR_NETWORK_ERROR:
1543                    if errno1 == errno.EADDRNOTAVAIL:
1544                        # Address not unavailable
1545                        if qh._client is not None:
1546                            raise SourceAddressBindError('Unable to bind to local address %s (%s)' % (qh._client, errno.errorcode[errno1]))
1547                        else:
1548                            raise SourceAddressBindError('Unable to bind to local address (%s)' % (errno.errorcode[errno1]))
1549                    elif errno1 == errno.EADDRINUSE or \
1550                            (errno1 == errno.EACCES and qtm.src is None):
1551                        # Address/port in use (EADDRINUSE) or insufficient
1552                        # permissions to bind to port
1553                        if qh.params['sport'] is not None:
1554                            raise PortBindError('Unable to bind to local port %d (%s)' % (qh.params['sport'], errno.errorcode[errno1]))
1555                        else:
1556                            raise PortBindError('Unable to bind to local port (%s)' % (errno.errorcode[errno1]))
1557                    elif qtm.src is None and errno1 not in (errno.EHOSTUNREACH, errno.ENETUNREACH, errno.EAFNOSUPPORT):
1558                        # If source is None it didn't bind properly.  If the
1559                        # errno1 value after bind() is EHOSTUNREACH or
1560                        # ENETUNREACH, it is because there was no proper IPv4
1561                        # or IPv6 connectivity (which is handled elsewhere).
1562                        # If socket() failed and resulted in an errno value of
1563                        # EAFNOSUPPORT, then likewise there is not IPv6
1564                        # support. In other cases, it was something unknown, so
1565                        # raise an error.
1566                        raise BindError('Unable to bind to local address (%s)' % (errno.errorcode.get(errno1, "unknown")))
1567
1568                # if src is None, then it is a connectivity issue on our
1569                # side, so don't record it in the responses
1570                if src is not None:
1571                    query.add_response(qh._server, src, response_obj, query.bailiwick)
1572
1573                # This query is now executed, at least in part
1574                query._executed = True
1575
1576            if newth.qtms:
1577                newth.init_req()
1578                bisect.insort(request_list, (query_time, newth))
1579
1580    def require_executed(func):
1581        def _func(self, *args, **kwargs):
1582            assert self._executed == True, "ExecutableDNSQuery has not been executed."
1583            return func(self, *args, **kwargs)
1584        return _func
1585
1586    def require_not_executed(func):
1587        def _func(self, *args, **kwargs):
1588            assert self._executed == False, "ExecutableDNSQuery has already been executed."
1589            return func(self, *args, **kwargs)
1590        return _func
1591
1592    def add_response(self, server, client, response, bailiwick):
1593        super(ExecutableDNSQuery, self).add_response(server, client, response, bailiwick)
1594        if not self.servers.difference(self.responses):
1595            self._executed = True
1596
1597    @require_not_executed
1598    def execute(self, ignore_queryid=True, tm=None, th_factories=None):
1599        self.execute_queries(self, ignore_queryid=ignore_queryid, tm=tm, th_factories=th_factories)
1600
1601    join = require_executed(DNSQuery.join)
1602    project = require_executed(DNSQuery.project)
1603    is_authoritative_answer_all = require_executed(DNSQuery.is_authoritative_answer_all)
1604    is_nxdomain_all = require_executed(DNSQuery.is_nxdomain_all)
1605    is_not_delegation_all = require_executed(DNSQuery.is_not_delegation_all)
1606    is_nxdomain_any = require_executed(DNSQuery.is_nxdomain_any)
1607
1608class DNSQueryFactory(object):
1609    '''A simple, extensible class interface for instantiating DNSQuery objects.'''
1610
1611    flags = 0
1612    edns = -1
1613    edns_max_udp_payload = 4096
1614    edns_flags = 0
1615    edns_options = []
1616
1617    tcp = False
1618
1619    query_timeout = 3.0
1620    max_attempts = 5
1621    lifetime = 15.0
1622
1623    response_handlers = []
1624
1625    def __new__(cls, qname, rdtype, rdclass, servers, bailiwick=None,
1626            client_ipv4=None, client_ipv6=None, port=53, odd_ports=None, cookie_jar=None, cookie_standin=None, cookie_bad=None,
1627            query_timeout=None, max_attempts=None, lifetime=None,
1628            executable=True):
1629
1630        if query_timeout is None:
1631            query_timeout = cls.query_timeout
1632        if max_attempts is None:
1633            max_attempts = cls.max_attempts
1634        if lifetime is None:
1635            lifetime = cls.lifetime
1636
1637        if executable:
1638            return ExecutableDNSQuery(qname, rdtype, rdclass, servers, bailiwick,
1639                client_ipv4, client_ipv6, port, odd_ports, cookie_jar, cookie_standin, cookie_bad,
1640                cls.flags, cls.edns, cls.edns_max_udp_payload, cls.edns_flags, cls.edns_options, cls.tcp,
1641                cls.response_handlers, query_timeout, max_attempts, lifetime)
1642
1643        else:
1644            return DNSQuery(qname, rdtype, rdclass,
1645                cls.flags, cls.edns, cls.edns_max_udp_payload, cls.edns_flags, cls.edns_options, cls.tcp)
1646
1647    def __init__(self, *args, **kwargs):
1648        raise NotImplemented()
1649
1650    @classmethod
1651    def add_mixin(cls, mixin_cls):
1652        class _foo(cls):
1653            flags = cls.flags | getattr(mixin_cls, 'flags', 0)
1654            edns_flags = cls.edns_flags | getattr(mixin_cls, 'edns_flags', 0)
1655            edns_options = cls.edns_options + copy.deepcopy(getattr(mixin_cls, 'edns_options', []))
1656        return _foo
1657
1658    @classmethod
1659    def get_cookie_opt(cls):
1660        try:
1661            return [o for o in cls.edns_options if o.otype == 10][0]
1662        except IndexError:
1663            return None
1664
1665    @classmethod
1666    def add_server_cookie(cls, server_cookie):
1667        cookie_opt = cls.get_cookie_opt()
1668        if cookie_opt is not None:
1669            if len(cookie_opt.data) != 8:
1670                raise TypeError('COOKIE option must have length of 8.')
1671            cookie_opt.data += server_cookie
1672        return cls
1673
1674    @classmethod
1675    def remove_cookie_option(cls):
1676        cookie_opt = cls.get_cookie_opt()
1677        if cookie_opt is not None:
1678            cls.edns_options.remove(cookie_opt)
1679        return cls
1680
1681class SimpleDNSQuery(DNSQueryFactory):
1682    '''A simple query, no frills.'''
1683
1684    pass
1685
1686class RecursiveDNSQuery(SimpleDNSQuery):
1687    '''A simple recursive query.'''
1688
1689    flags = SimpleDNSQuery.flags | dns.flags.RD
1690
1691class StandardQuery(SimpleDNSQuery):
1692    '''A standard old-school DNS query that handles truncated packets.'''
1693
1694    response_handlers = \
1695            SimpleDNSQuery.response_handlers + \
1696            [UseTCPOnTCFlagHandler()]
1697
1698class StandardRecursiveQuery(StandardQuery, RecursiveDNSQuery):
1699    '''A standard old-school recursive DNS query that handles truncated packets.'''
1700
1701    pass
1702
1703class StandardRecursiveQueryCD(StandardRecursiveQuery):
1704    '''A recursive DNS query that retries with checking disabled if the
1705    response code is SERVFAIL.'''
1706
1707    response_handlers = \
1708            StandardRecursiveQuery.response_handlers + \
1709            [SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL)]
1710
1711class EDNS0Query(StandardQuery):
1712    '''A standard query with EDNS0.'''
1713
1714    edns = 0
1715
1716class RecursiveEDNS0Query(EDNS0Query, RecursiveDNSQuery):
1717    '''A standard recursive query with EDNS0.'''
1718
1719    pass
1720
1721class DNSSECQuery(EDNS0Query):
1722    '''A standard query requesting DNSSEC records.'''
1723
1724    edns_flags = EDNS0Query.edns_flags | dns.flags.DO
1725
1726class RecursiveDNSSECQuery(DNSSECQuery, RecursiveDNSQuery):
1727    '''A standard recursive query requesting DNSSEC records.'''
1728
1729    pass
1730
1731class QuickDNSSECQuery(DNSSECQuery):
1732    '''A standard DNSSEC query, designed for quick turnaround.'''
1733
1734    response_handlers = DNSSECQuery.response_handlers + \
1735            [
1736                    AddServerCookieOnBADCOOKIE(),
1737                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1738                    DisableEDNSOnFormerrHandler(),
1739                    DisableEDNSOnRcodeHandler()
1740            ]
1741
1742    query_timeout = 1.0
1743    max_attempts = 1
1744    lifetime = 3.0
1745
1746class DiagnosticQuery(DNSSECQuery):
1747    '''A robust query with a number of handlers, designed to detect common DNS
1748    compatibility and connectivity issues.'''
1749
1750    response_handlers = DNSSECQuery.response_handlers + \
1751            [
1752                    AddServerCookieOnBADCOOKIE(),
1753                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1754                    DisableEDNSOnFormerrHandler(),
1755                    DisableEDNSOnRcodeHandler(),
1756                    ReduceUDPMaxPayloadOnTimeoutHandler(512, 4),
1757                    RemoveEDNSOptionOnTimeoutHandler(6),
1758                    ClearEDNSFlagOnTimeoutHandler(dns.flags.DO, 10),
1759                    DisableEDNSOnTimeoutHandler(11),
1760                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
1761                    ChangeTimeoutOnTimeoutHandler(1.0, 4),
1762                    ChangeTimeoutOnTimeoutHandler(2.0, 5),
1763                    ChangeTimeoutOnTimeoutHandler(1.0, 6),
1764            ]
1765    # For timeouts:
1766    #  1 - no change
1767    #  2 - change timeout to 2 seconds
1768    #  3 - no change
1769    #  4 - reduce udp max payload to 512; change timeout to 1 second
1770    #  5 - change timeout to 2 seconds
1771    #  6 - remove EDNS option (if any); change timeout to 1 second
1772    #  7 - remove EDNS option (if any)
1773    #  8 - remove EDNS option (if any)
1774    #  9 - remove EDNS option (if any)
1775    #  10 - clear DO flag;
1776    #  11 - disable EDNS
1777    #  12 - return (give up)
1778
1779    query_timeout = 1.0
1780    max_attempts = 12
1781    lifetime = 16.0
1782
1783class RecursiveDiagnosticQuery(RecursiveDNSSECQuery):
1784    '''A robust query to a cache with a number of handlers, designed to detect
1785    common DNS compatibility and connectivity issues.'''
1786
1787    response_handlers = DNSSECQuery.response_handlers + \
1788            [
1789                    AddServerCookieOnBADCOOKIE(),
1790                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1791                    DisableEDNSOnFormerrHandler(),
1792                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
1793                    DisableEDNSOnRcodeHandler(),
1794                    ReduceUDPMaxPayloadOnTimeoutHandler(512, 5),
1795                    RemoveEDNSOptionOnTimeoutHandler(7),
1796                    ClearEDNSFlagOnTimeoutHandler(dns.flags.DO, 11),
1797                    DisableEDNSOnTimeoutHandler(12),
1798                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
1799                    ChangeTimeoutOnTimeoutHandler(4.0, 3),
1800                    ChangeTimeoutOnTimeoutHandler(8.0, 4),
1801                    ChangeTimeoutOnTimeoutHandler(1.0, 5),
1802                    ChangeTimeoutOnTimeoutHandler(2.0, 6),
1803                    ChangeTimeoutOnTimeoutHandler(1.0, 7),
1804            ]
1805    # For timeouts:
1806    #  1 - no change
1807    #  2 - change timeout to 2 seconds
1808    #  3 - change timeout to 4 seconds
1809    #  4 - change timeout to 8 seconds
1810    #  5 - reduce udp max payload to 512; change timeout to 1 second
1811    #  6 - change timeout to 2 seconds
1812    #  7 - remove EDNS option (if any); change timeout to 1 second
1813    #  8 - remove EDNS option (if any)
1814    #  9 - remove EDNS option (if any)
1815    #  10 - remove EDNS option (if any)
1816    #  11 - clear DO flag
1817    #  12 - disable EDNS
1818    #  13 - return (give up)
1819
1820    query_timeout = 1.0
1821    max_attempts = 13
1822    lifetime = 26.0
1823
1824class TCPDiagnosticQuery(DNSSECQuery):
1825    '''A robust query with a number of handlers, designed to detect common DNS
1826    compatibility and connectivity issues over TCP.'''
1827
1828    tcp = True
1829
1830    response_handlers = \
1831            [
1832                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1833                    DisableEDNSOnFormerrHandler(),
1834                    DisableEDNSOnRcodeHandler(),
1835                    ChangeTimeoutOnTimeoutHandler(4.0, 2)
1836            ]
1837    # For timeouts:
1838    #  1 - no change
1839    #  2 - change timeout to 4 seconds
1840    #  3 - return
1841
1842    query_timeout = 2.0
1843    max_attempts = 3
1844    lifetime = 10.0
1845
1846class RecursiveTCPDiagnosticQuery(RecursiveDNSSECQuery):
1847    '''A robust query with a number of handlers, designed to detect common DNS
1848    compatibility and connectivity issues, beginning with TCP.'''
1849
1850    tcp = True
1851
1852    response_handlers = \
1853            [
1854                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1855                    DisableEDNSOnFormerrHandler(),
1856                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
1857                    DisableEDNSOnRcodeHandler(),
1858                    ChangeTimeoutOnTimeoutHandler(4.0, 2),
1859                    ChangeTimeoutOnTimeoutHandler(8.0, 3)
1860            ]
1861    # For timeouts:
1862    #  1 - no change
1863    #  2 - change timeout to 4 seconds
1864    #  3 - change timeout to 8 seconds
1865    #  4 - return
1866
1867    query_timeout = 2.0
1868    max_attempts = 4
1869    lifetime = 18.0
1870
1871class PMTUDiagnosticQuery(DNSSECQuery):
1872
1873    response_handlers = \
1874            [PMTUBoundingHandler(512, 4, 6, 1.0)] + \
1875            DNSSECQuery.response_handlers + \
1876            [
1877                    AddServerCookieOnBADCOOKIE(),
1878                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1879                    DisableEDNSOnFormerrHandler(),
1880                    DisableEDNSOnRcodeHandler(),
1881                    RemoveEDNSOptionOnTimeoutHandler(6),
1882                    ClearEDNSFlagOnTimeoutHandler(dns.flags.DO, 10),
1883                    DisableEDNSOnTimeoutHandler(11),
1884                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
1885                    ChangeTimeoutOnTimeoutHandler(1.0, 4),
1886                    ChangeTimeoutOnTimeoutHandler(2.0, 5),
1887                    ChangeTimeoutOnTimeoutHandler(1.0, 6),
1888            ]
1889    # For timeouts:
1890    #  1 - no change
1891    #  2 - change timeout to 2 seconds
1892    #  3 - no change
1893    #  4 - reduce udp max payload to 512; change timeout to 1 second
1894    #  5 - change timeout to 2 seconds
1895    #  6 - remove EDNS option (if any); change timeout to 1 second
1896    #  7 - remove EDNS option (if any)
1897    #  8 - remove EDNS option (if any)
1898    #  9 - remove EDNS option (if any)
1899    #  10 - clear DO flag;
1900    #  11 - disable EDNS
1901    #  12 - return (give up)
1902
1903    query_timeout = 1.0
1904    max_attempts = 12
1905    lifetime = 22.0 # set this a little longer due to pickle stage
1906
1907class RecursivePMTUDiagnosticQuery(RecursiveDNSSECQuery):
1908
1909    response_handlers = \
1910            [PMTUBoundingHandler(512, 5, 7, 1.0)] + \
1911            DNSSECQuery.response_handlers + \
1912            [
1913                    AddServerCookieOnBADCOOKIE(),
1914                    RemoveEDNSOptionOnRcodeHandler(dns.rcode.FORMERR),
1915                    DisableEDNSOnFormerrHandler(),
1916                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
1917                    DisableEDNSOnRcodeHandler(),
1918                    RemoveEDNSOptionOnTimeoutHandler(7),
1919                    ClearEDNSFlagOnTimeoutHandler(dns.flags.DO, 11),
1920                    DisableEDNSOnTimeoutHandler(12),
1921                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
1922                    ChangeTimeoutOnTimeoutHandler(4.0, 3),
1923                    ChangeTimeoutOnTimeoutHandler(8.0, 4),
1924                    ChangeTimeoutOnTimeoutHandler(1.0, 5),
1925                    ChangeTimeoutOnTimeoutHandler(2.0, 6),
1926                    ChangeTimeoutOnTimeoutHandler(1.0, 7),
1927            ]
1928    # For timeouts:
1929    #  1 - no change
1930    #  2 - change timeout to 2 seconds
1931    #  3 - change timeout to 4 seconds
1932    #  4 - change timeout to 8 seconds
1933    #  5 - reduce udp max payload to 512; change timeout to 1 second
1934    #  6 - change timeout to 2 seconds
1935    #  7 - remove EDNS option (if any); change timeout to 1 second
1936    #  8 - remove EDNS option (if any)
1937    #  9 - remove EDNS option (if any)
1938    #  10 - remove EDNS option (if any)
1939    #  11 - clear DO flag
1940    #  12 - disable EDNS
1941    #  13 - return (give up)
1942
1943    query_timeout = 1.0
1944    max_attempts = 13
1945    lifetime = 32.0 # set this a little longer due to pickle stage
1946
1947class TruncationDiagnosticQuery(DNSSECQuery):
1948    '''A simple query to test the results of a query with capabilities of only
1949    receiving back a small (512 byte) payload.'''
1950
1951    response_handlers = \
1952            [
1953                    AddServerCookieOnBADCOOKIE(),
1954                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
1955                    ChangeTimeoutOnTimeoutHandler(4.0, 3)
1956            ]
1957    # For timeouts:
1958    #  1 - no change
1959    #  2 - change timeout to 2 seconds
1960    #  3 - change timeout to 4 seconds
1961
1962    edns_max_udp_payload = 512
1963
1964    query_timeout = 1.0
1965    max_attempts = 4
1966    lifetime = 8.0
1967
1968class RecursiveTruncationDiagnosticQuery(DNSSECQuery, RecursiveDNSQuery):
1969    '''A simple recursive query to test the results of a query with
1970    capabilities of only receiving back a small (512 byte) payload.'''
1971
1972    response_handlers = \
1973            [
1974                    AddServerCookieOnBADCOOKIE(),
1975                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
1976                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
1977                    ChangeTimeoutOnTimeoutHandler(4.0, 3),
1978                    ChangeTimeoutOnTimeoutHandler(8.0, 4)
1979            ]
1980    # For timeouts:
1981    #  1 - no change
1982    #  2 - change timeout to 2 seconds
1983    #  3 - change timeout to 4 seconds
1984    #  4 - change timeout to 8 seconds
1985
1986    edns_max_udp_payload = 512
1987
1988    query_timeout = 1.0
1989    max_attempts = 5
1990    lifetime = 18.0
1991
1992class EDNSVersionDiagnosticQuery(SimpleDNSQuery):
1993    '''A query designed to test unknown EDNS version compatibility.'''
1994
1995    edns = 100
1996    edns_max_udp_payload = 512
1997
1998    response_handlers = \
1999            SimpleDNSQuery.response_handlers + \
2000            [
2001                    ChangeEDNSVersionOnTimeoutHandler(0, 4),
2002                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
2003                    ChangeTimeoutOnTimeoutHandler(1.0, 4)
2004            ]
2005    # For timeouts:
2006    #  1 - no change
2007    #  2 - change timeout to 2 seconds
2008    #  3 - no change
2009    #  4 - change EDNS version to 0; change timeout to 1 second
2010    #  5 - return
2011
2012    query_timeout = 1.0
2013    max_attempts = 5
2014    lifetime = 7.0
2015
2016class EDNSOptDiagnosticQuery(SimpleDNSQuery):
2017    '''A query designed to test unknown EDNS option compatibility.'''
2018
2019    edns = 0
2020    edns_max_udp_payload = 512
2021    edns_options = [dns.edns.GenericOption(100, b'')]
2022
2023    response_handlers = \
2024            SimpleDNSQuery.response_handlers + \
2025            [
2026                    AddServerCookieOnBADCOOKIE(),
2027                    RemoveEDNSOptionOnTimeoutHandler(4),
2028                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
2029                    ChangeTimeoutOnTimeoutHandler(1.0, 4)
2030            ]
2031
2032    # For timeouts:
2033    #  1 - no change
2034    #  2 - change timeout to 2 seconds
2035    #  3 - no change
2036    #  4 - remove EDNS option (if any); change timeout to 1 second
2037    #  5 - remove EDNS option (if any)
2038    #  6 - remove EDNS option (if any)
2039    #  7 - remove EDNS option (if any)
2040    #  8 - return
2041
2042    query_timeout = 1.0
2043    max_attempts = 8
2044    lifetime = 11.0
2045
2046class EDNSFlagDiagnosticQuery(SimpleDNSQuery):
2047    '''A query designed to test unknown EDNS flag compatibility.'''
2048
2049    edns = 0
2050    edns_max_udp_payload = 512
2051    edns_flags = SimpleDNSQuery.edns_flags | 0x80
2052
2053    response_handlers = \
2054            SimpleDNSQuery.response_handlers + \
2055            [
2056                    AddServerCookieOnBADCOOKIE(),
2057                    RemoveEDNSOptionOnTimeoutHandler(4),
2058                    ClearEDNSFlagOnTimeoutHandler(0x80, 8),
2059                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
2060                    ChangeTimeoutOnTimeoutHandler(1.0, 4)
2061            ]
2062
2063    # For timeouts:
2064    #  1 - no change
2065    #  2 - change timeout to 2 seconds
2066    #  3 - no change
2067    #  4 - remove EDNS option (if any); change timeout to 1 second
2068    #  5 - remove EDNS option (if any)
2069    #  6 - remove EDNS option (if any)
2070    #  7 - remove EDNS option (if any)
2071    #  8 - clear EDNS flag
2072    #  9 - return
2073
2074    query_timeout = 1.0
2075    max_attempts = 9
2076    lifetime = 12.0
2077
2078class RecursiveEDNSVersionDiagnosticQuery(SimpleDNSQuery):
2079    '''A query designed to test unknown EDNS version compatibility on recursive
2080    servers.'''
2081
2082    flags = dns.flags.RD
2083    edns = 100
2084    edns_max_udp_payload = 512
2085
2086    response_handlers = \
2087            SimpleDNSQuery.response_handlers + \
2088            [
2089                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
2090                    ChangeEDNSVersionOnTimeoutHandler(0, 5),
2091                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
2092                    ChangeTimeoutOnTimeoutHandler(4.0, 3),
2093                    ChangeTimeoutOnTimeoutHandler(8.0, 4),
2094                    ChangeTimeoutOnTimeoutHandler(1.0, 5)
2095            ]
2096    # For timeouts:
2097    #  1 - no change
2098    #  2 - change timeout to 2 seconds
2099    #  3 - change timeout to 4 seconds
2100    #  4 - change timeout to 8 seconds
2101    #  5 - change EDNS version to 0; change timeout to 1 second
2102    #  6 - return
2103
2104    query_timeout = 1.0
2105    max_attempts = 6
2106    lifetime = 18.0
2107
2108class RecursiveEDNSOptDiagnosticQuery(SimpleDNSQuery):
2109    '''A query designed to test unknown EDNS option compatibility on recursive
2110    servers.'''
2111
2112    flags = dns.flags.RD
2113    edns = 0
2114    edns_max_udp_payload = 512
2115    edns_options = [dns.edns.GenericOption(100, b'')]
2116
2117    response_handlers = \
2118            SimpleDNSQuery.response_handlers + \
2119            [
2120                    AddServerCookieOnBADCOOKIE(),
2121                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
2122                    RemoveEDNSOptionOnTimeoutHandler(5),
2123                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
2124                    ChangeTimeoutOnTimeoutHandler(4.0, 3),
2125                    ChangeTimeoutOnTimeoutHandler(8.0, 4),
2126                    ChangeTimeoutOnTimeoutHandler(1.0, 5)
2127            ]
2128
2129    # For timeouts:
2130    #  1 - no change
2131    #  2 - change timeout to 2 seconds
2132    #  3 - change timeout to 4 seconds
2133    #  4 - change timeout to 8 seconds
2134    #  5 - remove EDNS option (if any); change timeout to 1 second
2135    #  6 - remove EDNS option (if any)
2136    #  7 - remove EDNS option (if any)
2137    #  8 - remove EDNS option (if any)
2138    #  9 - return
2139
2140    query_timeout = 1.0
2141    max_attempts = 9
2142    lifetime = 21.0
2143
2144class RecursiveEDNSFlagDiagnosticQuery(SimpleDNSQuery):
2145    '''A query designed to test unknown EDNS flag compatibility on recursive
2146    servers.'''
2147
2148    flags = dns.flags.RD
2149    edns = 0
2150    edns_max_udp_payload = 512
2151    edns_flags = SimpleDNSQuery.edns_flags | 0x80
2152
2153    response_handlers = \
2154            SimpleDNSQuery.response_handlers + \
2155            [
2156                    AddServerCookieOnBADCOOKIE(),
2157                    SetFlagOnRcodeHandler(dns.flags.CD, dns.rcode.SERVFAIL),
2158                    RemoveEDNSOptionOnTimeoutHandler(5),
2159                    ClearEDNSFlagOnTimeoutHandler(0x80, 9),
2160                    ChangeTimeoutOnTimeoutHandler(2.0, 2),
2161                    ChangeTimeoutOnTimeoutHandler(4.0, 3),
2162                    ChangeTimeoutOnTimeoutHandler(8.0, 4),
2163                    ChangeTimeoutOnTimeoutHandler(1.0, 5)
2164            ]
2165
2166    # For timeouts:
2167    #  1 - no change
2168    #  2 - change timeout to 2 seconds
2169    #  3 - change timeout to 4 seconds
2170    #  4 - change timeout to 8 seconds
2171    #  5 - remove EDNS option (if any); change timeout to 1 second
2172    #  6 - remove EDNS option (if any)
2173    #  7 - remove EDNS option (if any)
2174    #  8 - remove EDNS option (if any)
2175    #  9 - clear EDNS flag
2176    #  10 - return
2177
2178    query_timeout = 1.0
2179    max_attempts = 10
2180    lifetime = 22.0
2181
2182def main():
2183    import json
2184    import sys
2185    import getopt
2186
2187    def usage():
2188        sys.stderr.write('Usage: %s [-r] [-j] <name> <type> <server> [<server>...]\n' % (sys.argv[0]))
2189        sys.exit(1)
2190
2191    try:
2192        opts, args = getopt.getopt(sys.argv[1:], 'rj')
2193        opts = dict(opts)
2194    except getopt.error:
2195        usage()
2196
2197    if len(args) < 3:
2198        usage()
2199
2200    if '-r' in opts:
2201        cls = RecursiveDiagnosticQuery
2202    else:
2203        cls = DiagnosticQuery
2204    d = cls(dns.name.from_text(args[0]), dns.rdatatype.from_text(args[1]), dns.rdataclass.IN, [IPAddr(x) for x in args[2:]])
2205    d.execute()
2206
2207    if '-j' in opts:
2208        print(json.dumps(d.serialize(), indent=4, separators=(',', ': ')))
2209    else:
2210        print('Responses for %s/%s:' % (args[0], args[1]))
2211        for server in d.responses:
2212            for client, response in d.responses[server].items():
2213                if response.message is not None:
2214                    print('   from %s: %s (%d bytes in %dms)' % (server, repr(response.message), len(response.message.to_wire()), int(response.response_time*1000)))
2215                else:
2216                    print('   from %s: (ERR: %s) (%dms)' % (server, repr(response.error), int(response.response_time*1000)))
2217
2218                print('   (src: %s)' % (client))
2219                if response.history:
2220                    print('       (history: %s)' % (response.history))
2221
2222if __name__ == '__main__':
2223    main()
2224