source: trunk/packages/invirt-dns/invirt-dns @ 1490

Last change on this file since 1490 was 1490, checked in by broder, 15 years ago

Add support for basic quoting in the DNS zone file

  • Property svn:executable set to *
File size: 8.0 KB
Line 
1#!/usr/bin/python
2from twisted.internet import reactor
3from twisted.names import server
4from twisted.names import dns
5from twisted.names import common
6from twisted.names import authority
7from twisted.internet import defer
8from twisted.python import failure
9
10from invirt.config import structs as config
11import invirt.database
12import psycopg2
13import sqlalchemy
14import time
15import re
16
17class DatabaseAuthority(common.ResolverBase):
18    """An Authority that is loaded from a file."""
19
20    soa = None
21
22    def __init__(self, domains=None, database=None):
23        common.ResolverBase.__init__(self)
24        if database is not None:
25            invirt.database.connect(database)
26        else:
27            invirt.database.connect()
28        if domains is not None:
29            self.domains = domains
30        else:
31            self.domains = config.dns.domains
32        ns = config.dns.nameservers[0]
33        self.soa = dns.Record_SOA(mname=ns.hostname,
34                                  rname=config.dns.contact.replace('@','.',1),
35                                  serial=1, refresh=3600, retry=900,
36                                  expire=3600000, minimum=21600, ttl=3600)
37        self.ns = dns.Record_NS(name=ns.hostname, ttl=3600)
38        record = dns.Record_A(address=ns.ip, ttl=3600)
39        self.ns1 = dns.RRHeader(ns.hostname, dns.A, dns.IN,
40                                3600, record, auth=True)
41
42   
43    def _lookup(self, name, cls, type, timeout = None):
44        for i in range(3):
45            try:
46                value = self._lookup_unsafe(name, cls, type, timeout = None)
47            except (psycopg2.OperationalError, sqlalchemy.exceptions.SQLError):
48                if i == 2:
49                    raise
50                print "Reloading database"
51                time.sleep(0.5)
52                continue
53            else:
54                return value
55
56    def _lookup_unsafe(self, name, cls, type, timeout):
57        invirt.database.clear_cache()
58       
59        ttl = 900
60        name = name.lower()
61
62        if name in self.domains:
63            domain = name
64        else:
65            # Look for the longest-matching domain.  (This works because domain
66            # will remain bound after breaking out of the loop.)
67            best_domain = ''
68            for domain in self.domains:
69                if name.endswith('.'+domain) and len(domain) > len(best_domain):
70                    best_domain = domain
71            if best_domain == '':
72                return defer.fail(failure.Failure(dns.DomainError(name)))
73            domain = best_domain
74        results = []
75        authority = []
76        additional = [self.ns1]
77        authority.append(dns.RRHeader(domain, dns.NS, dns.IN,
78                                      3600, self.ns, auth=True))
79
80        if cls == dns.IN:
81            host = name[:-len(domain)-1]
82            if not host: # Request for the domain itself.
83                if type in (dns.A, dns.ALL_RECORDS):
84                    record = dns.Record_A(config.dns.nameservers[0].ip, ttl)
85                    results.append(dns.RRHeader(name, dns.A, dns.IN, 
86                                                ttl, record, auth=True))
87                elif type == dns.NS:
88                    results.append(dns.RRHeader(domain, dns.NS, dns.IN,
89                                                ttl, self.ns, auth=True))
90                    authority = []
91                elif type == dns.SOA:
92                    results.append(dns.RRHeader(domain, dns.SOA, dns.IN,
93                                                ttl, self.soa, auth=True))
94            else: # Request for a subdomain.
95                value = invirt.database.Machine.query().filter_by(name=host).first()
96                if value is None or not value.nics:
97                    return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
98                ip = value.nics[0].ip
99                if ip is None:  #Deactivated?
100                    return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
101
102                if type in (dns.A, dns.ALL_RECORDS):
103                    record = dns.Record_A(ip, ttl)
104                    results.append(dns.RRHeader(name, dns.A, dns.IN, 
105                                                ttl, record, auth=True))
106                elif type == dns.SOA:
107                    results.append(dns.RRHeader(domain, dns.SOA, dns.IN,
108                                                ttl, self.soa, auth=True))
109            if len(results) == 0:
110                authority = []
111                additional = []
112            return defer.succeed((results, authority, additional))
113        else:
114            #Doesn't exist
115            return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
116
117class QuotingBindAuthority(authority.BindAuthority):
118    """
119    A BindAuthority that (almost) deals with quoting correctly
120   
121    This will catch double quotes as marking the start or end of a
122    quoted phrase, unless the double quote is escaped by a backslash
123    """
124    # Grab everything up to the first whitespace character or
125    # quotation mark not proceeded by a backslash
126    whitespace_re = re.compile(r'(.*?)([\t\n\x0b\x0c\r ]+|(?<!\\)")')
127    def collapseContinuations(self, lines):
128        L = []
129        state = 0
130        for line in lines:
131            if state == 0:
132                if line.find('(') == -1:
133                    L.append(line)
134                else:
135                    L.append(line[:line.find('(')])
136                    state = 1
137            else:
138                if line.find(')') != -1:
139                    L[-1] += ' ' + line[:line.find(')')]
140                    state = 0
141                else:
142                    L[-1] += ' ' + line
143        lines = L
144        L = []
145        for line in lines:
146            in_quote = False
147            split_line = []
148            while len(line) > 0:
149                match = self.whitespace_re.match(line)
150                if match is None:
151                    # If there's no match, that means that there's no
152                    # whitespace in the rest of the line, so it should
153                    # be treated as a single entity, quoted or not
154                    #
155                    # This also means that a closing quote isn't
156                    # strictly necessary if the line ends the quote
157                    substr = line
158                    end = ''
159                else:
160                    substr, end = match.groups()
161               
162                if in_quote:
163                    # If we're in the middle of the quote, the string
164                    # we just grabbed belongs at the end of the
165                    # previous string
166                    #
167                    # Including the whitespace! Unless it's not
168                    # whitespace and is actually a closequote instead
169                    split_line[-1] += substr + (end if end != '"' else '')
170                else:
171                    # If we're not in the middle of a quote, than this
172                    # is the next new string
173                    split_line.append(substr)
174               
175                if end == '"':
176                    in_quote = not in_quote
177               
178                # Then strip off what we just processed
179                line = line[len(substr + end):]
180            L.append(split_line)
181        return filter(None, L)
182
183if '__main__' == __name__:
184    resolvers = []
185    for zone in config.dns.zone_files:
186        for origin in config.dns.domains:
187            r = QuotingBindAuthority(zone)
188            # This sucks, but if I want a generic zone file, I have to
189            # reload the information by hand
190            r.origin = origin
191            lines = open(zone).readlines()
192            lines = r.collapseContinuations(r.stripComments(lines))
193            r.parseLines(lines)
194           
195            resolvers.append(r)
196    resolvers.append(DatabaseAuthority())
197
198    verbosity = 0
199    f = server.DNSServerFactory(authorities=resolvers, verbose=verbosity)
200    p = dns.DNSDatagramProtocol(f)
201    f.noisy = p.noisy = verbosity
202   
203    reactor.listenUDP(53, p)
204    reactor.listenTCP(53, f)
205    reactor.run()
Note: See TracBrowser for help on using the repository browser.