1# Copyright (C) 2003-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
16import socket
17import struct
18
19import dns.ipv4
20import dns.rdata
21
22_proto_tcp = socket.getprotobyname('tcp')
23_proto_udp = socket.getprotobyname('udp')
24
25class WKS(dns.rdata.Rdata):
26    """WKS record
27
28    @ivar address: the address
29    @type address: string
30    @ivar protocol: the protocol
31    @type protocol: int
32    @ivar bitmap: the bitmap
33    @type bitmap: string
34    @see: RFC 1035"""
35
36    __slots__ = ['address', 'protocol', 'bitmap']
37
38    def __init__(self, rdclass, rdtype, address, protocol, bitmap):
39        super(WKS, self).__init__(rdclass, rdtype)
40        self.address = address
41        self.protocol = protocol
42        self.bitmap = bitmap
43
44    def to_text(self, origin=None, relativize=True, **kw):
45        bits = []
46        for i in xrange(0, len(self.bitmap)):
47            byte = ord(self.bitmap[i])
48            for j in xrange(0, 8):
49                if byte & (0x80 >> j):
50                    bits.append(str(i * 8 + j))
51        text = ' '.join(bits)
52        return '%s %d %s' % (self.address, self.protocol, text)
53
54    def from_text(cls, rdclass, rdtype, tok, origin = None, relativize = True):
55        address = tok.get_string()
56        protocol = tok.get_string()
57        if protocol.isdigit():
58            protocol = int(protocol)
59        else:
60            protocol = socket.getprotobyname(protocol)
61        bitmap = []
62        while 1:
63            token = tok.get().unescape()
64            if token.is_eol_or_eof():
65                break
66            if token.value.isdigit():
67                serv = int(token.value)
68            else:
69                if protocol != _proto_udp and protocol != _proto_tcp:
70                    raise NotImplementedError("protocol must be TCP or UDP")
71                if protocol == _proto_udp:
72                    protocol_text = "udp"
73                else:
74                    protocol_text = "tcp"
75                serv = socket.getservbyname(token.value, protocol_text)
76            i = serv // 8
77            l = len(bitmap)
78            if l < i + 1:
79                for j in xrange(l, i + 1):
80                    bitmap.append('\x00')
81            bitmap[i] = chr(ord(bitmap[i]) | (0x80 >> (serv % 8)))
82        bitmap = dns.rdata._truncate_bitmap(bitmap)
83        return cls(rdclass, rdtype, address, protocol, bitmap)
84
85    from_text = classmethod(from_text)
86
87    def to_wire(self, file, compress = None, origin = None):
88        file.write(dns.ipv4.inet_aton(self.address))
89        protocol = struct.pack('!B', self.protocol)
90        file.write(protocol)
91        file.write(self.bitmap)
92
93    def from_wire(cls, rdclass, rdtype, wire, current, rdlen, origin = None):
94        address = dns.ipv4.inet_ntoa(wire[current : current + 4])
95        protocol, = struct.unpack('!B', wire[current + 4 : current + 5])
96        current += 5
97        rdlen -= 5
98        bitmap = wire[current : current + rdlen]
99        return cls(rdclass, rdtype, address, protocol, bitmap)
100
101    from_wire = classmethod(from_wire)
102
103    def _cmp(self, other):
104        sa = dns.ipv4.inet_aton(self.address)
105        oa = dns.ipv4.inet_aton(other.address)
106        v = cmp(sa, oa)
107        if v == 0:
108            sp = struct.pack('!B', self.protocol)
109            op = struct.pack('!B', other.protocol)
110            v = cmp(sp, op)
111            if v == 0:
112                v = cmp(self.bitmap, other.bitmap)
113        return v
114