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.internet import defer |
---|
7 | from twisted.python import failure |
---|
8 | |
---|
9 | from invirt.config import structs as config |
---|
10 | import invirt.database |
---|
11 | import psycopg2 |
---|
12 | import sqlalchemy |
---|
13 | import time |
---|
14 | |
---|
15 | class DatabaseAuthority(common.ResolverBase): |
---|
16 | """An Authority that is loaded from a file.""" |
---|
17 | |
---|
18 | soa = None |
---|
19 | |
---|
20 | def __init__(self, domains=None, database=None): |
---|
21 | common.ResolverBase.__init__(self) |
---|
22 | if database is not None: |
---|
23 | invirt.database.connect(database) |
---|
24 | else: |
---|
25 | invirt.database.connect() |
---|
26 | if domains is not None: |
---|
27 | self.domains = domains |
---|
28 | else: |
---|
29 | self.domains = config.dns.domains |
---|
30 | ns = config.dns.nameservers[0] |
---|
31 | self.soa = dns.Record_SOA(mname=ns.hostname, |
---|
32 | rname=config.dns.contact.replace('@','.',1), |
---|
33 | serial=1, refresh=3600, retry=900, |
---|
34 | expire=3600000, minimum=21600, ttl=3600) |
---|
35 | self.ns = dns.Record_NS(name=ns.hostname, ttl=3600) |
---|
36 | record = dns.Record_A(address=ns.ip, ttl=3600) |
---|
37 | self.ns1 = dns.RRHeader(ns.hostname, dns.A, dns.IN, |
---|
38 | 3600, record, auth=True) |
---|
39 | |
---|
40 | |
---|
41 | def _lookup(self, name, cls, type, timeout = None): |
---|
42 | for i in range(3): |
---|
43 | try: |
---|
44 | value = self._lookup_unsafe(name, cls, type, timeout = None) |
---|
45 | except (psycopg2.OperationalError, sqlalchemy.exceptions.SQLError): |
---|
46 | if i == 2: |
---|
47 | raise |
---|
48 | print "Reloading database" |
---|
49 | time.sleep(0.5) |
---|
50 | continue |
---|
51 | else: |
---|
52 | return value |
---|
53 | |
---|
54 | def _lookup_unsafe(self, name, cls, type, timeout): |
---|
55 | invirt.database.clear_cache() |
---|
56 | |
---|
57 | ttl = 900 |
---|
58 | name = name.lower() |
---|
59 | |
---|
60 | # XXX hack for the transition to two separate dev/prod clusters |
---|
61 | if 'dev.xvm.mit.edu' in self.domains and name.endswith('prod.xvm.mit.edu'): |
---|
62 | # import time, sys |
---|
63 | # print time.localtime(), 'handling prod request', name |
---|
64 | # sys.stdout.flush() |
---|
65 | |
---|
66 | # Point the client in the right direction for prod requests. |
---|
67 | authority = dns.RRHeader('prod.xvm.mit.edu', dns.NS, dns.IN, 3600, |
---|
68 | dns.Record_NS(name='ns1.prod.xvm.mit.edu', ttl=3600), auth=True) |
---|
69 | additional = dns.RRHeader('ns1.prod.xvm.mit.edu', dns.A, dns.IN, 3600, |
---|
70 | dns.Record_A(address='18.181.0.221', ttl=3600), auth=True) |
---|
71 | return defer.succeed(([], [authority], [additional])) |
---|
72 | |
---|
73 | if name in self.domains: |
---|
74 | domain = name |
---|
75 | else: |
---|
76 | # Look for the longest-matching domain. (This works because domain |
---|
77 | # will remain bound after breaking out of the loop.) |
---|
78 | best_domain = '' |
---|
79 | for domain in self.domains: |
---|
80 | if name.endswith('.'+domain) and len(domain) > len(best_domain): |
---|
81 | best_domain = domain |
---|
82 | if best_domain == '': |
---|
83 | return defer.fail(failure.Failure(dns.DomainError(name))) |
---|
84 | domain = best_domain |
---|
85 | results = [] |
---|
86 | authority = [] |
---|
87 | additional = [self.ns1] |
---|
88 | authority.append(dns.RRHeader(domain, dns.NS, dns.IN, |
---|
89 | 3600, self.ns, auth=True)) |
---|
90 | |
---|
91 | if cls == dns.IN: |
---|
92 | host = name[:-len(domain)-1] |
---|
93 | if not host: # Request for the domain itself. |
---|
94 | if type in (dns.A, dns.ALL_RECORDS): |
---|
95 | record = dns.Record_A(config.dns.nameservers[0].ip, ttl) |
---|
96 | results.append(dns.RRHeader(name, dns.A, dns.IN, |
---|
97 | ttl, record, auth=True)) |
---|
98 | elif type == dns.NS: |
---|
99 | results.append(dns.RRHeader(domain, dns.NS, dns.IN, |
---|
100 | ttl, self.ns, auth=True)) |
---|
101 | authority = [] |
---|
102 | elif type == dns.SOA: |
---|
103 | results.append(dns.RRHeader(domain, dns.SOA, dns.IN, |
---|
104 | ttl, self.soa, auth=True)) |
---|
105 | else: # Request for a subdomain. |
---|
106 | if 'passup' in dir(config.dns) and host in config.dns.passup: |
---|
107 | record = dns.Record_CNAME('%s.%s' % (host, config.dns.parent), ttl) |
---|
108 | return defer.succeed(( |
---|
109 | [dns.RRHeader(name, dns.CNAME, dns.IN, ttl, record, auth=True)], |
---|
110 | [], [])) |
---|
111 | |
---|
112 | value = invirt.database.Machine.query().filter_by(name=host).one() |
---|
113 | if value is None or not value.nics: |
---|
114 | return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name))) |
---|
115 | ip = value.nics[0].ip |
---|
116 | if ip is None: #Deactivated? |
---|
117 | return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name))) |
---|
118 | |
---|
119 | if type in (dns.A, dns.ALL_RECORDS): |
---|
120 | record = dns.Record_A(ip, ttl) |
---|
121 | results.append(dns.RRHeader(name, dns.A, dns.IN, |
---|
122 | ttl, record, auth=True)) |
---|
123 | elif type == dns.SOA: |
---|
124 | results.append(dns.RRHeader(domain, dns.SOA, dns.IN, |
---|
125 | ttl, self.soa, auth=True)) |
---|
126 | if len(results) == 0: |
---|
127 | authority = [] |
---|
128 | additional = [] |
---|
129 | return defer.succeed((results, authority, additional)) |
---|
130 | else: |
---|
131 | #Doesn't exist |
---|
132 | return defer.fail(failure.Failure(dns.AuthoritativeDomainError(name))) |
---|
133 | |
---|
134 | if '__main__' == __name__: |
---|
135 | resolver = DatabaseAuthority() |
---|
136 | |
---|
137 | verbosity = 0 |
---|
138 | f = server.DNSServerFactory(authorities=[resolver], verbose=verbosity) |
---|
139 | p = dns.DNSDatagramProtocol(f) |
---|
140 | f.noisy = p.noisy = verbosity |
---|
141 | |
---|
142 | reactor.listenUDP(53, p) |
---|
143 | reactor.listenTCP(53, f) |
---|
144 | reactor.run() |
---|