1# Copyright (C) 2001-2007, 2009, 2010 Nominum, Inc.
2#
3# Permission to use, copy, modify, and distribute this software and its
4# documentation for any purpose with or without fee is hereby granted,
5# provided that the above copyright notice and this permission notice
6# appear in all copies.
7#
8# THE SOFTWARE IS PROVIDED "AS IS" AND NOMINUM DISCLAIMS ALL WARRANTIES
9# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL NOMINUM BE LIABLE FOR
11# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
14# OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
15
16"""DNS Messages"""
17
18import cStringIO
19import random
20import struct
21import sys
22import time
23
24import dns.exception
25import dns.flags
26import dns.name
27import dns.opcode
28import dns.entropy
29import dns.rcode
30import dns.rdata
31import dns.rdataclass
32import dns.rdatatype
33import dns.rrset
34import dns.renderer
35import dns.tsig
36
37class ShortHeader(dns.exception.FormError):
38    """Raised if the DNS packet passed to from_wire() is too short."""
39    pass
40
41class TrailingJunk(dns.exception.FormError):
42    """Raised if the DNS packet passed to from_wire() has extra junk
43    at the end of it."""
44    pass
45
46class UnknownHeaderField(dns.exception.DNSException):
47    """Raised if a header field name is not recognized when converting from
48    text into a message."""
49    pass
50
51class BadEDNS(dns.exception.FormError):
52    """Raised if an OPT record occurs somewhere other than the start of
53    the additional data section."""
54    pass
55
56class BadTSIG(dns.exception.FormError):
57    """Raised if a TSIG record occurs somewhere other than the end of
58    the additional data section."""
59    pass
60
61class UnknownTSIGKey(dns.exception.DNSException):
62    """Raised if we got a TSIG but don't know the key."""
63    pass
64
65class Message(object):
66    """A DNS message.
67
68    @ivar id: The query id; the default is a randomly chosen id.
69    @type id: int
70    @ivar flags: The DNS flags of the message.  @see: RFC 1035 for an
71    explanation of these flags.
72    @type flags: int
73    @ivar question: The question section.
74    @type question: list of dns.rrset.RRset objects
75    @ivar answer: The answer section.
76    @type answer: list of dns.rrset.RRset objects
77    @ivar authority: The authority section.
78    @type authority: list of dns.rrset.RRset objects
79    @ivar additional: The additional data section.
80    @type additional: list of dns.rrset.RRset objects
81    @ivar edns: The EDNS level to use.  The default is -1, no Edns.
82    @type edns: int
83    @ivar ednsflags: The EDNS flags
84    @type ednsflags: long
85    @ivar payload: The EDNS payload size.  The default is 0.
86    @type payload: int
87    @ivar options: The EDNS options
88    @type options: list of dns.edns.Option objects
89    @ivar request_payload: The associated request's EDNS payload size.
90    @type request_payload: int
91    @ivar keyring: The TSIG keyring to use.  The default is None.
92    @type keyring: dict
93    @ivar keyname: The TSIG keyname to use.  The default is None.
94    @type keyname: dns.name.Name object
95    @ivar keyalgorithm: The TSIG key algorithm to use.  The default is
96    dns.tsig.default_algorithm.
97    @type keyalgorithm: string
98    @ivar request_mac: The TSIG MAC of the request message associated with
99    this message; used when validating TSIG signatures.   @see: RFC 2845 for
100    more information on TSIG fields.
101    @type request_mac: string
102    @ivar fudge: TSIG time fudge; default is 300 seconds.
103    @type fudge: int
104    @ivar original_id: TSIG original id; defaults to the message's id
105    @type original_id: int
106    @ivar tsig_error: TSIG error code; default is 0.
107    @type tsig_error: int
108    @ivar other_data: TSIG other data.
109    @type other_data: string
110    @ivar mac: The TSIG MAC for this message.
111    @type mac: string
112    @ivar xfr: Is the message being used to contain the results of a DNS
113    zone transfer?  The default is False.
114    @type xfr: bool
115    @ivar origin: The origin of the zone in messages which are used for
116    zone transfers or for DNS dynamic updates.  The default is None.
117    @type origin: dns.name.Name object
118    @ivar tsig_ctx: The TSIG signature context associated with this
119    message.  The default is None.
120    @type tsig_ctx: hmac.HMAC object
121    @ivar had_tsig: Did the message decoded from wire format have a TSIG
122    signature?
123    @type had_tsig: bool
124    @ivar multi: Is this message part of a multi-message sequence?  The
125    default is false.  This variable is used when validating TSIG signatures
126    on messages which are part of a zone transfer.
127    @type multi: bool
128    @ivar first: Is this message standalone, or the first of a multi
129    message sequence?  This variable is used when validating TSIG signatures
130    on messages which are part of a zone transfer.
131    @type first: bool
132    @ivar index: An index of rrsets in the message.  The index key is
133    (section, name, rdclass, rdtype, covers, deleting).  Indexing can be
134    disabled by setting the index to None.
135    @type index: dict
136    """
137
138    def __init__(self, id=None):
139        if id is None:
140            self.id = dns.entropy.random_16()
141        else:
142            self.id = id
143        self.flags = 0
144        self.question = []
145        self.answer = []
146        self.authority = []
147        self.additional = []
148        self.edns = -1
149        self.ednsflags = 0
150        self.payload = 0
151        self.options = []
152        self.request_payload = 0
153        self.keyring = None
154        self.keyname = None
155        self.keyalgorithm = dns.tsig.default_algorithm
156        self.request_mac = ''
157        self.other_data = ''
158        self.tsig_error = 0
159        self.fudge = 300
160        self.original_id = self.id
161        self.mac = ''
162        self.xfr = False
163        self.origin = None
164        self.tsig_ctx = None
165        self.had_tsig = False
166        self.multi = False
167        self.first = True
168        self.index = {}
169
170    def __repr__(self):
171        return '<DNS message, ID ' + `self.id` + '>'
172
173    def __str__(self):
174        return self.to_text()
175
176    def to_text(self,  origin=None, relativize=True, **kw):
177        """Convert the message to text.
178
179        The I{origin}, I{relativize}, and any other keyword
180        arguments are passed to the rrset to_wire() method.
181
182        @rtype: string
183        """
184
185        s = cStringIO.StringIO()
186        print >> s, 'id %d' % self.id
187        print >> s, 'opcode %s' % \
188              dns.opcode.to_text(dns.opcode.from_flags(self.flags))
189        rc = dns.rcode.from_flags(self.flags, self.ednsflags)
190        print >> s, 'rcode %s' % dns.rcode.to_text(rc)
191        print >> s, 'flags %s' % dns.flags.to_text(self.flags)
192        if self.edns >= 0:
193            print >> s, 'edns %s' % self.edns
194            if self.ednsflags != 0:
195                print >> s, 'eflags %s' % \
196                      dns.flags.edns_to_text(self.ednsflags)
197            print >> s, 'payload', self.payload
198        is_update = dns.opcode.is_update(self.flags)
199        if is_update:
200            print >> s, ';ZONE'
201        else:
202            print >> s, ';QUESTION'
203        for rrset in self.question:
204            print >> s, rrset.to_text(origin, relativize, **kw)
205        if is_update:
206            print >> s, ';PREREQ'
207        else:
208            print >> s, ';ANSWER'
209        for rrset in self.answer:
210            print >> s, rrset.to_text(origin, relativize, **kw)
211        if is_update:
212            print >> s, ';UPDATE'
213        else:
214            print >> s, ';AUTHORITY'
215        for rrset in self.authority:
216            print >> s, rrset.to_text(origin, relativize, **kw)
217        print >> s, ';ADDITIONAL'
218        for rrset in self.additional:
219            print >> s, rrset.to_text(origin, relativize, **kw)
220        #
221        # We strip off the final \n so the caller can print the result without
222        # doing weird things to get around eccentricities in Python print
223        # formatting
224        #
225        return s.getvalue()[:-1]
226
227    def __eq__(self, other):
228        """Two messages are equal if they have the same content in the
229        header, question, answer, and authority sections.
230        @rtype: bool"""
231        if not isinstance(other, Message):
232            return False
233        if self.id != other.id:
234            return False
235        if self.flags != other.flags:
236            return False
237        for n in self.question:
238            if n not in other.question:
239                return False
240        for n in other.question:
241            if n not in self.question:
242                return False
243        for n in self.answer:
244            if n not in other.answer:
245                return False
246        for n in other.answer:
247            if n not in self.answer:
248                return False
249        for n in self.authority:
250            if n not in other.authority:
251                return False
252        for n in other.authority:
253            if n not in self.authority:
254                return False
255        return True
256
257    def __ne__(self, other):
258        """Are two messages not equal?
259        @rtype: bool"""
260        return not self.__eq__(other)
261
262    def is_response(self, other):
263        """Is other a response to self?
264        @rtype: bool"""
265        if other.flags & dns.flags.QR == 0 or \
266           self.id != other.id or \
267           dns.opcode.from_flags(self.flags) != \
268           dns.opcode.from_flags(other.flags):
269            return False
270        if dns.rcode.from_flags(other.flags, other.ednsflags) != \
271               dns.rcode.NOERROR:
272            return True
273        if dns.opcode.is_update(self.flags):
274            return True
275        for n in self.question:
276            if n not in other.question:
277                return False
278        for n in other.question:
279            if n not in self.question:
280                return False
281        return True
282
283    def section_number(self, section):
284        if section is self.question:
285            return 0
286        elif section is self.answer:
287            return 1
288        elif section is self.authority:
289            return 2
290        elif section is self.additional:
291            return 3
292        else:
293            raise ValueError('unknown section')
294
295    def find_rrset(self, section, name, rdclass, rdtype,
296                   covers=dns.rdatatype.NONE, deleting=None, create=False,
297                   force_unique=False):
298        """Find the RRset with the given attributes in the specified section.
299
300        @param section: the section of the message to look in, e.g.
301        self.answer.
302        @type section: list of dns.rrset.RRset objects
303        @param name: the name of the RRset
304        @type name: dns.name.Name object
305        @param rdclass: the class of the RRset
306        @type rdclass: int
307        @param rdtype: the type of the RRset
308        @type rdtype: int
309        @param covers: the covers value of the RRset
310        @type covers: int
311        @param deleting: the deleting value of the RRset
312        @type deleting: int
313        @param create: If True, create the RRset if it is not found.
314        The created RRset is appended to I{section}.
315        @type create: bool
316        @param force_unique: If True and create is also True, create a
317        new RRset regardless of whether a matching RRset exists already.
318        @type force_unique: bool
319        @raises KeyError: the RRset was not found and create was False
320        @rtype: dns.rrset.RRset object"""
321
322        key = (self.section_number(section),
323               name, rdclass, rdtype, covers, deleting)
324        if not force_unique:
325            if not self.index is None:
326                rrset = self.index.get(key)
327                if not rrset is None:
328                    return rrset
329            else:
330                for rrset in section:
331                    if rrset.match(name, rdclass, rdtype, covers, deleting):
332                        return rrset
333        if not create:
334            raise KeyError
335        rrset = dns.rrset.RRset(name, rdclass, rdtype, covers, deleting)
336        section.append(rrset)
337        if not self.index is None:
338            self.index[key] = rrset
339        return rrset
340
341    def get_rrset(self, section, name, rdclass, rdtype,
342                  covers=dns.rdatatype.NONE, deleting=None, create=False,
343                  force_unique=False):
344        """Get the RRset with the given attributes in the specified section.
345
346        If the RRset is not found, None is returned.
347
348        @param section: the section of the message to look in, e.g.
349        self.answer.
350        @type section: list of dns.rrset.RRset objects
351        @param name: the name of the RRset
352        @type name: dns.name.Name object
353        @param rdclass: the class of the RRset
354        @type rdclass: int
355        @param rdtype: the type of the RRset
356        @type rdtype: int
357        @param covers: the covers value of the RRset
358        @type covers: int
359        @param deleting: the deleting value of the RRset
360        @type deleting: int
361        @param create: If True, create the RRset if it is not found.
362        The created RRset is appended to I{section}.
363        @type create: bool
364        @param force_unique: If True and create is also True, create a
365        new RRset regardless of whether a matching RRset exists already.
366        @type force_unique: bool
367        @rtype: dns.rrset.RRset object or None"""
368
369        try:
370            rrset = self.find_rrset(section, name, rdclass, rdtype, covers,
371                                    deleting, create, force_unique)
372        except KeyError:
373            rrset = None
374        return rrset
375
376    def to_wire(self, origin=None, max_size=0, **kw):
377        """Return a string containing the message in DNS compressed wire
378        format.
379
380        Additional keyword arguments are passed to the rrset to_wire()
381        method.
382
383        @param origin: The origin to be appended to any relative names.
384        @type origin: dns.name.Name object
385        @param max_size: The maximum size of the wire format output; default
386        is 0, which means 'the message's request payload, if nonzero, or
387        65536'.
388        @type max_size: int
389        @raises dns.exception.TooBig: max_size was exceeded
390        @rtype: string
391        """
392
393        if max_size == 0:
394            if self.request_payload != 0:
395                max_size = self.request_payload
396            else:
397                max_size = 65535
398        if max_size < 512:
399            max_size = 512
400        elif max_size > 65535:
401            max_size = 65535
402        r = dns.renderer.Renderer(self.id, self.flags, max_size, origin)
403        for rrset in self.question:
404            r.add_question(rrset.name, rrset.rdtype, rrset.rdclass)
405        for rrset in self.answer:
406            r.add_rrset(dns.renderer.ANSWER, rrset, **kw)
407        for rrset in self.authority:
408            r.add_rrset(dns.renderer.AUTHORITY, rrset, **kw)
409        if self.edns >= 0:
410            r.add_edns(self.edns, self.ednsflags, self.payload, self.options)
411        for rrset in self.additional:
412            r.add_rrset(dns.renderer.ADDITIONAL, rrset, **kw)
413        r.write_header()
414        if not self.keyname is None:
415            r.add_tsig(self.keyname, self.keyring[self.keyname],
416                       self.fudge, self.original_id, self.tsig_error,
417                       self.other_data, self.request_mac,
418                       self.keyalgorithm)
419            self.mac = r.mac
420        return r.get_wire()
421
422    def use_tsig(self, keyring, keyname=None, fudge=300,
423                 original_id=None, tsig_error=0, other_data='',
424                 algorithm=dns.tsig.default_algorithm):
425        """When sending, a TSIG signature using the specified keyring
426        and keyname should be added.
427
428        @param keyring: The TSIG keyring to use; defaults to None.
429        @type keyring: dict
430        @param keyname: The name of the TSIG key to use; defaults to None.
431        The key must be defined in the keyring.  If a keyring is specified
432        but a keyname is not, then the key used will be the first key in the
433        keyring.  Note that the order of keys in a dictionary is not defined,
434        so applications should supply a keyname when a keyring is used, unless
435        they know the keyring contains only one key.
436        @type keyname: dns.name.Name or string
437        @param fudge: TSIG time fudge; default is 300 seconds.
438        @type fudge: int
439        @param original_id: TSIG original id; defaults to the message's id
440        @type original_id: int
441        @param tsig_error: TSIG error code; default is 0.
442        @type tsig_error: int
443        @param other_data: TSIG other data.
444        @type other_data: string
445        @param algorithm: The TSIG algorithm to use; defaults to
446        dns.tsig.default_algorithm
447        """
448
449        self.keyring = keyring
450        if keyname is None:
451            self.keyname = self.keyring.keys()[0]
452        else:
453            if isinstance(keyname, (str, unicode)):
454                keyname = dns.name.from_text(keyname)
455            self.keyname = keyname
456        self.keyalgorithm = algorithm
457        self.fudge = fudge
458        if original_id is None:
459            self.original_id = self.id
460        else:
461            self.original_id = original_id
462        self.tsig_error = tsig_error
463        self.other_data = other_data
464
465    def use_edns(self, edns=0, ednsflags=0, payload=1280, request_payload=None, options=None):
466        """Configure EDNS behavior.
467        @param edns: The EDNS level to use.  Specifying None, False, or -1
468        means 'do not use EDNS', and in this case the other parameters are
469        ignored.  Specifying True is equivalent to specifying 0, i.e. 'use
470        EDNS0'.
471        @type edns: int or bool or None
472        @param ednsflags: EDNS flag values.
473        @type ednsflags: int
474        @param payload: The EDNS sender's payload field, which is the maximum
475        size of UDP datagram the sender can handle.
476        @type payload: int
477        @param request_payload: The EDNS payload size to use when sending
478        this message.  If not specified, defaults to the value of payload.
479        @type request_payload: int or None
480        @param options: The EDNS options
481        @type options: None or list of dns.edns.Option objects
482        @see: RFC 2671
483        """
484        if edns is None or edns is False:
485            edns = -1
486        if edns is True:
487            edns = 0
488        if request_payload is None:
489            request_payload = payload
490        if edns < 0:
491            ednsflags = 0
492            payload = 0
493            request_payload = 0
494            options = []
495        else:
496            # make sure the EDNS version in ednsflags agrees with edns
497            ednsflags &= 0xFF00FFFFL
498            ednsflags |= (edns << 16)
499            if options is None:
500                options = []
501        self.edns = edns
502        self.ednsflags = ednsflags
503        self.payload = payload
504        self.options = options
505        self.request_payload = request_payload
506
507    def want_dnssec(self, wanted=True):
508        """Enable or disable 'DNSSEC desired' flag in requests.
509        @param wanted: Is DNSSEC desired?  If True, EDNS is enabled if
510        required, and then the DO bit is set.  If False, the DO bit is
511        cleared if EDNS is enabled.
512        @type wanted: bool
513        """
514        if wanted:
515            if self.edns < 0:
516                self.use_edns()
517            self.ednsflags |= dns.flags.DO
518        elif self.edns >= 0:
519            self.ednsflags &= ~dns.flags.DO
520
521    def rcode(self):
522        """Return the rcode.
523        @rtype: int
524        """
525        return dns.rcode.from_flags(self.flags, self.ednsflags)
526
527    def set_rcode(self, rcode):
528        """Set the rcode.
529        @param rcode: the rcode
530        @type rcode: int
531        """
532        (value, evalue) = dns.rcode.to_flags(rcode)
533        self.flags &= 0xFFF0
534        self.flags |= value
535        self.ednsflags &= 0x00FFFFFFL
536        self.ednsflags |= evalue
537        if self.ednsflags != 0 and self.edns < 0:
538            self.edns = 0
539
540    def opcode(self):
541        """Return the opcode.
542        @rtype: int
543        """
544        return dns.opcode.from_flags(self.flags)
545
546    def set_opcode(self, opcode):
547        """Set the opcode.
548        @param opcode: the opcode
549        @type opcode: int
550        """
551        self.flags &= 0x87FF
552        self.flags |= dns.opcode.to_flags(opcode)
553
554class _WireReader(object):
555    """Wire format reader.
556
557    @ivar wire: the wire-format message.
558    @type wire: string
559    @ivar message: The message object being built
560    @type message: dns.message.Message object
561    @ivar current: When building a message object from wire format, this
562    variable contains the offset from the beginning of wire of the next octet
563    to be read.
564    @type current: int
565    @ivar updating: Is the message a dynamic update?
566    @type updating: bool
567    @ivar one_rr_per_rrset: Put each RR into its own RRset?
568    @type one_rr_per_rrset: bool
569    @ivar zone_rdclass: The class of the zone in messages which are
570    DNS dynamic updates.
571    @type zone_rdclass: int
572    """
573
574    def __init__(self, wire, message, question_only=False,
575                 one_rr_per_rrset=False):
576        self.wire = wire
577        self.message = message
578        self.current = 0
579        self.updating = False
580        self.zone_rdclass = dns.rdataclass.IN
581        self.question_only = question_only
582        self.one_rr_per_rrset = one_rr_per_rrset
583
584    def _get_question(self, qcount):
585        """Read the next I{qcount} records from the wire data and add them to
586        the question section.
587        @param qcount: the number of questions in the message
588        @type qcount: int"""
589
590        if self.updating and qcount > 1:
591            raise dns.exception.FormError
592
593        for i in xrange(0, qcount):
594            (qname, used) = dns.name.from_wire(self.wire, self.current)
595            if not self.message.origin is None:
596                qname = qname.relativize(self.message.origin)
597            self.current = self.current + used
598            (rdtype, rdclass) = \
599                     struct.unpack('!HH',
600                                   self.wire[self.current:self.current + 4])
601            self.current = self.current + 4
602            self.message.find_rrset(self.message.question, qname,
603                                    rdclass, rdtype, create=True,
604                                    force_unique=True)
605            if self.updating:
606                self.zone_rdclass = rdclass
607
608    def _get_section(self, section, count):
609        """Read the next I{count} records from the wire data and add them to
610        the specified section.
611        @param section: the section of the message to which to add records
612        @type section: list of dns.rrset.RRset objects
613        @param count: the number of records to read
614        @type count: int"""
615
616        if self.updating or self.one_rr_per_rrset:
617            force_unique = True
618        else:
619            force_unique = False
620        seen_opt = False
621        for i in xrange(0, count):
622            rr_start = self.current
623            (name, used) = dns.name.from_wire(self.wire, self.current)
624            absolute_name = name
625            if not self.message.origin is None:
626                name = name.relativize(self.message.origin)
627            self.current = self.current + used
628            (rdtype, rdclass, ttl, rdlen) = \
629                     struct.unpack('!HHIH',
630                                   self.wire[self.current:self.current + 10])
631            self.current = self.current + 10
632            if rdtype == dns.rdatatype.OPT:
633                if not section is self.message.additional or seen_opt:
634                    raise BadEDNS
635                self.message.payload = rdclass
636                self.message.ednsflags = ttl
637                self.message.edns = (ttl & 0xff0000) >> 16
638                self.message.options = []
639                current = self.current
640                optslen = rdlen
641                while optslen > 0:
642                    (otype, olen) = \
643                            struct.unpack('!HH',
644                                          self.wire[current:current + 4])
645                    current = current + 4
646                    opt = dns.edns.option_from_wire(otype, self.wire, current, olen)
647                    self.message.options.append(opt)
648                    current = current + olen
649                    optslen = optslen - 4 - olen
650                seen_opt = True
651            elif rdtype == dns.rdatatype.TSIG:
652                if not (section is self.message.additional and
653                        i == (count - 1)):
654                    raise BadTSIG
655                if self.message.keyring is None:
656                    raise UnknownTSIGKey('got signed message without keyring')
657                secret = self.message.keyring.get(absolute_name)
658                if secret is None:
659                    raise UnknownTSIGKey("key '%s' unknown" % name)
660                self.message.tsig_ctx = \
661                                      dns.tsig.validate(self.wire,
662                                          absolute_name,
663                                          secret,
664                                          int(time.time()),
665                                          self.message.request_mac,
666                                          rr_start,
667                                          self.current,
668                                          rdlen,
669                                          self.message.tsig_ctx,
670                                          self.message.multi,
671                                          self.message.first)
672                self.message.had_tsig = True
673            else:
674                if ttl < 0:
675                    ttl = 0
676                if self.updating and \
677                   (rdclass == dns.rdataclass.ANY or
678                    rdclass == dns.rdataclass.NONE):
679                    deleting = rdclass
680                    rdclass = self.zone_rdclass
681                else:
682                    deleting = None
683                if deleting == dns.rdataclass.ANY or \
684                   (deleting == dns.rdataclass.NONE and \
685                    section == self.message.answer):
686                    covers = dns.rdatatype.NONE
687                    rd = None
688                else:
689                    rd = dns.rdata.from_wire(rdclass, rdtype, self.wire,
690                                             self.current, rdlen,
691                                             self.message.origin)
692                    covers = rd.covers()
693                if self.message.xfr and rdtype == dns.rdatatype.SOA:
694                    force_unique = True
695                rrset = self.message.find_rrset(section, name,
696                                                rdclass, rdtype, covers,
697                                                deleting, True, force_unique)
698                if not rd is None:
699                    rrset.add(rd, ttl)
700            self.current = self.current + rdlen
701
702    def read(self):
703        """Read a wire format DNS message and build a dns.message.Message
704        object."""
705
706        l = len(self.wire)
707        if l < 12:
708            raise ShortHeader
709        (self.message.id, self.message.flags, qcount, ancount,
710         aucount, adcount) = struct.unpack('!HHHHHH', self.wire[:12])
711        self.current = 12
712        if dns.opcode.is_update(self.message.flags):
713            self.updating = True
714        self._get_question(qcount)
715        if self.question_only:
716            return
717        self._get_section(self.message.answer, ancount)
718        self._get_section(self.message.authority, aucount)
719        self._get_section(self.message.additional, adcount)
720        if self.current != l:
721            raise TrailingJunk
722        if self.message.multi and self.message.tsig_ctx and \
723               not self.message.had_tsig:
724            self.message.tsig_ctx.update(self.wire)
725
726
727def from_wire(wire, keyring=None, request_mac='', xfr=False, origin=None,
728              tsig_ctx = None, multi = False, first = True,
729              question_only = False, one_rr_per_rrset = False):
730    """Convert a DNS wire format message into a message
731    object.
732
733    @param keyring: The keyring to use if the message is signed.
734    @type keyring: dict
735    @param request_mac: If the message is a response to a TSIG-signed request,
736    I{request_mac} should be set to the MAC of that request.
737    @type request_mac: string
738    @param xfr: Is this message part of a zone transfer?
739    @type xfr: bool
740    @param origin: If the message is part of a zone transfer, I{origin}
741    should be the origin name of the zone.
742    @type origin: dns.name.Name object
743    @param tsig_ctx: The ongoing TSIG context, used when validating zone
744    transfers.
745    @type tsig_ctx: hmac.HMAC object
746    @param multi: Is this message part of a multiple message sequence?
747    @type multi: bool
748    @param first: Is this message standalone, or the first of a multi
749    message sequence?
750    @type first: bool
751    @param question_only: Read only up to the end of the question section?
752    @type question_only: bool
753    @param one_rr_per_rrset: Put each RR into its own RRset
754    @type one_rr_per_rrset: bool
755    @raises ShortHeader: The message is less than 12 octets long.
756    @raises TrailingJunk: There were octets in the message past the end
757    of the proper DNS message.
758    @raises BadEDNS: An OPT record was in the wrong section, or occurred more
759    than once.
760    @raises BadTSIG: A TSIG record was not the last record of the additional
761    data section.
762    @rtype: dns.message.Message object"""
763
764    m = Message(id=0)
765    m.keyring = keyring
766    m.request_mac = request_mac
767    m.xfr = xfr
768    m.origin = origin
769    m.tsig_ctx = tsig_ctx
770    m.multi = multi
771    m.first = first
772
773    reader = _WireReader(wire, m, question_only, one_rr_per_rrset)
774    reader.read()
775
776    return m
777
778
779class _TextReader(object):
780    """Text format reader.
781
782    @ivar tok: the tokenizer
783    @type tok: dns.tokenizer.Tokenizer object
784    @ivar message: The message object being built
785    @type message: dns.message.Message object
786    @ivar updating: Is the message a dynamic update?
787    @type updating: bool
788    @ivar zone_rdclass: The class of the zone in messages which are
789    DNS dynamic updates.
790    @type zone_rdclass: int
791    @ivar last_name: The most recently read name when building a message object
792    from text format.
793    @type last_name: dns.name.Name object
794    """
795
796    def __init__(self, text, message):
797        self.message = message
798        self.tok = dns.tokenizer.Tokenizer(text)
799        self.last_name = None
800        self.zone_rdclass = dns.rdataclass.IN
801        self.updating = False
802
803    def _header_line(self, section):
804        """Process one line from the text format header section."""
805
806        token = self.tok.get()
807        what = token.value
808        if what == 'id':
809            self.message.id = self.tok.get_int()
810        elif what == 'flags':
811            while True:
812                token = self.tok.get()
813                if not token.is_identifier():
814                    self.tok.unget(token)
815                    break
816                self.message.flags = self.message.flags | \
817                                     dns.flags.from_text(token.value)
818            if dns.opcode.is_update(self.message.flags):
819                self.updating = True
820        elif what == 'edns':
821            self.message.edns = self.tok.get_int()
822            self.message.ednsflags = self.message.ednsflags | \
823                                     (self.message.edns << 16)
824        elif what == 'eflags':
825            if self.message.edns < 0:
826                self.message.edns = 0
827            while True:
828                token = self.tok.get()
829                if not token.is_identifier():
830                    self.tok.unget(token)
831                    break
832                self.message.ednsflags = self.message.ednsflags | \
833                              dns.flags.edns_from_text(token.value)
834        elif what == 'payload':
835            self.message.payload = self.tok.get_int()
836            if self.message.edns < 0:
837                self.message.edns = 0
838        elif what == 'opcode':
839            text = self.tok.get_string()
840            self.message.flags = self.message.flags | \
841                      dns.opcode.to_flags(dns.opcode.from_text(text))
842        elif what == 'rcode':
843            text = self.tok.get_string()
844            self.message.set_rcode(dns.rcode.from_text(text))
845        else:
846            raise UnknownHeaderField
847        self.tok.get_eol()
848
849    def _question_line(self, section):
850        """Process one line from the text format question section."""
851
852        token = self.tok.get(want_leading = True)
853        if not token.is_whitespace():
854            self.last_name = dns.name.from_text(token.value, None)
855        name = self.last_name
856        token = self.tok.get()
857        if not token.is_identifier():
858            raise dns.exception.SyntaxError
859        # Class
860        try:
861            rdclass = dns.rdataclass.from_text(token.value)
862            token = self.tok.get()
863            if not token.is_identifier():
864                raise dns.exception.SyntaxError
865        except dns.exception.SyntaxError:
866            raise dns.exception.SyntaxError
867        except:
868            rdclass = dns.rdataclass.IN
869        # Type
870        rdtype = dns.rdatatype.from_text(token.value)
871        self.message.find_rrset(self.message.question, name,
872                                rdclass, rdtype, create=True,
873                                force_unique=True)
874        if self.updating:
875            self.zone_rdclass = rdclass
876        self.tok.get_eol()
877
878    def _rr_line(self, section):
879        """Process one line from the text format answer, authority, or
880        additional data sections.
881        """
882
883        deleting = None
884        # Name
885        token = self.tok.get(want_leading = True)
886        if not token.is_whitespace():
887            self.last_name = dns.name.from_text(token.value, None)
888        name = self.last_name
889        token = self.tok.get()
890        if not token.is_identifier():
891            raise dns.exception.SyntaxError
892        # TTL
893        try:
894            ttl = int(token.value, 0)
895            token = self.tok.get()
896            if not token.is_identifier():
897                raise dns.exception.SyntaxError
898        except dns.exception.SyntaxError:
899            raise dns.exception.SyntaxError
900        except:
901            ttl = 0
902        # Class
903        try:
904            rdclass = dns.rdataclass.from_text(token.value)
905            token = self.tok.get()
906            if not token.is_identifier():
907                raise dns.exception.SyntaxError
908            if rdclass == dns.rdataclass.ANY or rdclass == dns.rdataclass.NONE:
909                deleting = rdclass
910                rdclass = self.zone_rdclass
911        except dns.exception.SyntaxError:
912            raise dns.exception.SyntaxError
913        except:
914            rdclass = dns.rdataclass.IN
915        # Type
916        rdtype = dns.rdatatype.from_text(token.value)
917        token = self.tok.get()
918        if not token.is_eol_or_eof():
919            self.tok.unget(token)
920            rd = dns.rdata.from_text(rdclass, rdtype, self.tok, None)
921            covers = rd.covers()
922        else:
923            rd = None
924            covers = dns.rdatatype.NONE
925        rrset = self.message.find_rrset(section, name,
926                                        rdclass, rdtype, covers,
927                                        deleting, True, self.updating)
928        if not rd is None:
929            rrset.add(rd, ttl)
930
931    def read(self):
932        """Read a text format DNS message and build a dns.message.Message
933        object."""
934
935        line_method = self._header_line
936        section = None
937        while 1:
938            token = self.tok.get(True, True)
939            if token.is_eol_or_eof():
940                break
941            if token.is_comment():
942                u = token.value.upper()
943                if u == 'HEADER':
944                    line_method = self._header_line
945                elif u == 'QUESTION' or u == 'ZONE':
946                    line_method = self._question_line
947                    section = self.message.question
948                elif u == 'ANSWER' or u == 'PREREQ':
949                    line_method = self._rr_line
950                    section = self.message.answer
951                elif u == 'AUTHORITY' or u == 'UPDATE':
952                    line_method = self._rr_line
953                    section = self.message.authority
954                elif u == 'ADDITIONAL':
955                    line_method = self._rr_line
956                    section = self.message.additional
957                self.tok.get_eol()
958                continue
959            self.tok.unget(token)
960            line_method(section)
961
962
963def from_text(text):
964    """Convert the text format message into a message object.
965
966    @param text: The text format message.
967    @type text: string
968    @raises UnknownHeaderField:
969    @raises dns.exception.SyntaxError:
970    @rtype: dns.message.Message object"""
971
972    # 'text' can also be a file, but we don't publish that fact
973    # since it's an implementation detail.  The official file
974    # interface is from_file().
975
976    m = Message()
977
978    reader = _TextReader(text, m)
979    reader.read()
980
981    return m
982
983def from_file(f):
984    """Read the next text format message from the specified file.
985
986    @param f: file or string.  If I{f} is a string, it is treated
987    as the name of a file to open.
988    @raises UnknownHeaderField:
989    @raises dns.exception.SyntaxError:
990    @rtype: dns.message.Message object"""
991
992    if sys.hexversion >= 0x02030000:
993        # allow Unicode filenames; turn on universal newline support
994        str_type = basestring
995        opts = 'rU'
996    else:
997        str_type = str
998        opts = 'r'
999    if isinstance(f, str_type):
1000        f = file(f, opts)
1001        want_close = True
1002    else:
1003        want_close = False
1004
1005    try:
1006        m = from_text(f)
1007    finally:
1008        if want_close:
1009            f.close()
1010    return m
1011
1012def make_query(qname, rdtype, rdclass = dns.rdataclass.IN, use_edns=None,
1013               want_dnssec=False):
1014    """Make a query message.
1015
1016    The query name, type, and class may all be specified either
1017    as objects of the appropriate type, or as strings.
1018
1019    The query will have a randomly choosen query id, and its DNS flags
1020    will be set to dns.flags.RD.
1021
1022    @param qname: The query name.
1023    @type qname: dns.name.Name object or string
1024    @param rdtype: The desired rdata type.
1025    @type rdtype: int
1026    @param rdclass: The desired rdata class; the default is class IN.
1027    @type rdclass: int
1028    @param use_edns: The EDNS level to use; the default is None (no EDNS).
1029    See the description of dns.message.Message.use_edns() for the possible
1030    values for use_edns and their meanings.
1031    @type use_edns: int or bool or None
1032    @param want_dnssec: Should the query indicate that DNSSEC is desired?
1033    @type want_dnssec: bool
1034    @rtype: dns.message.Message object"""
1035
1036    if isinstance(qname, (str, unicode)):
1037        qname = dns.name.from_text(qname)
1038    if isinstance(rdtype, str):
1039        rdtype = dns.rdatatype.from_text(rdtype)
1040    if isinstance(rdclass, str):
1041        rdclass = dns.rdataclass.from_text(rdclass)
1042    m = Message()
1043    m.flags |= dns.flags.RD
1044    m.find_rrset(m.question, qname, rdclass, rdtype, create=True,
1045                 force_unique=True)
1046    m.use_edns(use_edns)
1047    m.want_dnssec(want_dnssec)
1048    return m
1049
1050def make_response(query, recursion_available=False, our_payload=8192):
1051    """Make a message which is a response for the specified query.
1052    The message returned is really a response skeleton; it has all
1053    of the infrastructure required of a response, but none of the
1054    content.
1055
1056    The response's question section is a shallow copy of the query's
1057    question section, so the query's question RRsets should not be
1058    changed.
1059
1060    @param query: the query to respond to
1061    @type query: dns.message.Message object
1062    @param recursion_available: should RA be set in the response?
1063    @type recursion_available: bool
1064    @param our_payload: payload size to advertise in EDNS responses; default
1065    is 8192.
1066    @type our_payload: int
1067    @rtype: dns.message.Message object"""
1068
1069    if query.flags & dns.flags.QR:
1070        raise dns.exception.FormError('specified query message is not a query')
1071    response = dns.message.Message(query.id)
1072    response.flags = dns.flags.QR | (query.flags & dns.flags.RD)
1073    if recursion_available:
1074        response.flags |= dns.flags.RA
1075    response.set_opcode(query.opcode())
1076    response.question = list(query.question)
1077    if query.edns >= 0:
1078        response.use_edns(0, 0, our_payload, query.payload)
1079    if not query.keyname is None:
1080        response.keyname = query.keyname
1081        response.keyring = query.keyring
1082        response.request_mac = query.mac
1083    return response
1084