Some quick touch-ups in Yang's quoting code - and a changelog entry
[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     # Match either a quoted or unquoted string literal followed by
125     # whitespace or the end of line.  This yields two groups, one of
126     # which has a match, and the other of which is None, depending on
127     # whether the string literal was quoted or unquoted; this is what
128     # necessitates the subsequent filtering out of groups that are
129     # None.
130     string_pat = \
131             re.compile(r'"((?:[^"\\]|\\.)*)"|((?:[^\\\s]|\\.)+)(?:\s+|\s*$)')
132
133     # For interpreting escapes.
134     escape_pat = re.compile(r'\\(.)')
135
136     def collapseContinuations(self, lines):
137         L = []
138         state = 0
139         for line in lines:
140             if state == 0:
141                 if line.find('(') == -1:
142                     L.append(line)
143                 else:
144                     L.append(line[:line.find('(')])
145                     state = 1
146             else:
147                 if line.find(')') != -1:
148                     L[-1] += ' ' + line[:line.find(')')]
149                     state = 0
150                 else:
151                     L[-1] += ' ' + line
152         lines = L
153         L = []
154
155         for line in lines:
156             in_quote = False
157             split_line = []
158             for m in self.string_pat.finditer(line):
159                 [x] = [x for x in m.groups() if x is not None]
160                 split_line.append(self.escape_pat.sub(r'\1', x))
161             L.append(split_line)
162         return filter(None, L)
163
164 if '__main__' == __name__:
165     resolvers = []
166     for zone in config.dns.zone_files:
167         for origin in config.dns.domains:
168             r = QuotingBindAuthority(zone)
169             # This sucks, but if I want a generic zone file, I have to
170             # reload the information by hand
171             r.origin = origin
172             lines = open(zone).readlines()
173             lines = r.collapseContinuations(r.stripComments(lines))
174             r.parseLines(lines)
175             
176             resolvers.append(r)
177     resolvers.append(DatabaseAuthority())
178
179     verbosity = 0
180     f = server.DNSServerFactory(authorities=resolvers, verbose=verbosity)
181     p = dns.DNSDatagramProtocol(f)
182     f.noisy = p.noisy = verbosity
183     
184     reactor.listenUDP(53, p)
185     reactor.listenTCP(53, f)
186     reactor.run()