Add support for basic quoting in the DNS zone file
[invirt/packages/invirt-dns.git] / invirt-dns
1 #!/usr/bin/python
2 from twisted.internet import reactor
3 from twisted.names import server
4 from twisted.names import dns
5 from twisted.names import common
6 from twisted.names import authority
7 from twisted.internet import defer
8 from twisted.python import failure
9
10 from invirt.config import structs as config
11 import invirt.database
12 import psycopg2
13 import sqlalchemy
14 import time
15 import re
16
17 class 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
117 class 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
183 if '__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()