Reverse-resolution support in invirt-dns
[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.common import InvirtConfigError
11 from invirt.config import structs as config
12 import invirt.database
13 import psycopg2
14 import sqlalchemy
15 import time
16 import re
17
18 class DatabaseAuthority(common.ResolverBase):
19     """An Authority that is loaded from a file."""
20
21     soa = None
22
23     def __init__(self, domains=None, database=None):
24         common.ResolverBase.__init__(self)
25         if database is not None:
26             invirt.database.connect(database)
27         else:
28             invirt.database.connect()
29         if domains is not None:
30             self.domains = domains
31         else:
32             self.domains = config.dns.domains
33         ns = config.dns.nameservers[0]
34         self.soa = dns.Record_SOA(mname=ns.hostname,
35                                   rname=config.dns.contact.replace('@','.',1),
36                                   serial=1, refresh=3600, retry=900,
37                                   expire=3600000, minimum=21600, ttl=3600)
38         self.ns = dns.Record_NS(name=ns.hostname, ttl=3600)
39         record = dns.Record_A(address=ns.ip, ttl=3600)
40         self.ns1 = dns.RRHeader(ns.hostname, dns.A, dns.IN,
41                                 3600, record, auth=True)
42
43     
44     def _lookup(self, name, cls, type, timeout = None):
45         for i in range(3):
46             try:
47                 value = self._lookup_unsafe(name, cls, type, timeout = None)
48             except (psycopg2.OperationalError, sqlalchemy.exceptions.SQLError):
49                 if i == 2:
50                     raise
51                 print "Reloading database"
52                 time.sleep(0.5)
53                 continue
54             else:
55                 return value
56
57     def _lookup_unsafe(self, name, cls, type, timeout):
58         invirt.database.clear_cache()
59         
60         ttl = 900
61         name = name.lower()
62
63         if name in self.domains:
64             domain = name
65         else:
66             # Look for the longest-matching domain.  (This works because domain
67             # will remain bound after breaking out of the loop.)
68             best_domain = ''
69             for domain in self.domains:
70                 if name.endswith('.'+domain) and len(domain) > len(best_domain):
71                     best_domain = domain
72             if best_domain == '':
73                 return defer.fail(failure.Failure(dns.DomainError(name)))
74             domain = best_domain
75         results = []
76         authority = []
77         additional = [self.ns1]
78         authority.append(dns.RRHeader(domain, dns.NS, dns.IN,
79                                       3600, self.ns, auth=True))
80
81         if cls == dns.IN:
82             host = name[:-len(domain)-1]
83             if not host: # Request for the domain itself.
84                 if type in (dns.A, dns.ALL_RECORDS):
85                     record = dns.Record_A(config.dns.nameservers[0].ip, ttl)
86                     results.append(dns.RRHeader(name, dns.A, dns.IN, 
87                                                 ttl, record, auth=True))
88                 elif type == dns.NS:
89                     results.append(dns.RRHeader(domain, dns.NS, dns.IN,
90                                                 ttl, self.ns, auth=True))
91                     authority = []
92                 elif type == dns.SOA:
93                     results.append(dns.RRHeader(domain, dns.SOA, dns.IN,
94                                                 ttl, self.soa, auth=True))
95             else: # Request for a subdomain.
96                 if name.endswith(".in-addr.arpa"): # Reverse resolution here
97                     if type in (dns.PTR, dns.ALL_RECORDS):
98                         ip = '.'.join(reversed(name.split('.')[:-2]))
99                         value = invirt.database.NIC.query.filter_by(ip=ip).first()
100                         if value and value.hostname:
101                             hostname = value.hostname
102                             if '.' not in hostname:
103                                 hostname = hostname + "." + config.dns.domains[0]
104                             record = dns.Record_PTR(hostname, ttl)
105                             results.append(dns.RRHeader(name, dns.PTR, dns.IN,
106                                                         ttl, record, auth=True))
107                         else: # IP address doesn't point to an active host
108                             return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
109                     # FIXME: Should only return success with no records if the name actually exists
110                 else: # Forward resolution here
111                     value = invirt.database.NIC.query.filter_by(hostname=host).first()
112                     if value:
113                         ip = value.ip
114                     else:
115                         value = invirt.database.Machine.query().filter_by(name=host).first()
116                         if value:
117                             ip = value.nics[0].ip
118                         else:
119                             return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
120                 
121                     if ip is None:
122                         return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
123
124                     if type in (dns.A, dns.ALL_RECORDS):
125                         record = dns.Record_A(ip, ttl)
126                         results.append(dns.RRHeader(name, dns.A, dns.IN, 
127                                                     ttl, record, auth=True))
128                     elif type == dns.SOA:
129                         results.append(dns.RRHeader(domain, dns.SOA, dns.IN,
130                                                     ttl, self.soa, auth=True))
131             if len(results) == 0:
132                 authority = []
133                 additional = []
134             return defer.succeed((results, authority, additional))
135         else:
136             #Doesn't exist
137             return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name)))
138
139 class QuotingBindAuthority(authority.BindAuthority):
140     """
141     A BindAuthority that (almost) deals with quoting correctly
142     
143     This will catch double quotes as marking the start or end of a
144     quoted phrase, unless the double quote is escaped by a backslash
145     """
146     # Match either a quoted or unquoted string literal followed by
147     # whitespace or the end of line.  This yields two groups, one of
148     # which has a match, and the other of which is None, depending on
149     # whether the string literal was quoted or unquoted; this is what
150     # necessitates the subsequent filtering out of groups that are
151     # None.
152     string_pat = \
153             re.compile(r'"((?:[^"\\]|\\.)*)"|((?:[^\\\s]|\\.)+)(?:\s+|\s*$)')
154
155     # For interpreting escapes.
156     escape_pat = re.compile(r'\\(.)')
157
158     def collapseContinuations(self, lines):
159         L = []
160         state = 0
161         for line in lines:
162             if state == 0:
163                 if line.find('(') == -1:
164                     L.append(line)
165                 else:
166                     L.append(line[:line.find('(')])
167                     state = 1
168             else:
169                 if line.find(')') != -1:
170                     L[-1] += ' ' + line[:line.find(')')]
171                     state = 0
172                 else:
173                     L[-1] += ' ' + line
174         lines = L
175         L = []
176
177         for line in lines:
178             in_quote = False
179             split_line = []
180             for m in self.string_pat.finditer(line):
181                 [x] = [x for x in m.groups() if x is not None]
182                 split_line.append(self.escape_pat.sub(r'\1', x))
183             L.append(split_line)
184         return filter(None, L)
185
186 if '__main__' == __name__:
187     resolvers = []
188     try:
189         for zone in config.dns.zone_files:
190             for origin in config.dns.domains:
191                 r = QuotingBindAuthority(zone)
192                 # This sucks, but if I want a generic zone file, I have to
193                 # reload the information by hand
194                 r.origin = origin
195                 lines = open(zone).readlines()
196                 lines = r.collapseContinuations(r.stripComments(lines))
197                 r.parseLines(lines)
198                 
199                 resolvers.append(r)
200     except InvirtConfigError:
201         # Don't care if zone_files isn't defined
202         pass
203     resolvers.append(DatabaseAuthority())
204
205     verbosity = 0
206     f = server.DNSServerFactory(authorities=resolvers, verbose=verbosity)
207     p = dns.DNSDatagramProtocol(f)
208     f.noisy = p.noisy = verbosity
209     
210     reactor.listenUDP(53, p)
211     reactor.listenTCP(53, f)
212     reactor.run()