# -*- coding: utf-8 -*- import logging import random import re from twisted.internet import defer from twisted.names import dns, server, client, common, resolve from twisted.python import failure from . import sdns logger = logging.getLogger(__name__) typeToMethod = { dns.A: 'lookupAddress', dns.AAAA: 'lookupIPV6Address', dns.A6: 'lookupAddress6', dns.NS: 'lookupNameservers', dns.CNAME: 'lookupCanonicalName', dns.SOA: 'lookupAuthority', dns.MB: 'lookupMailBox', dns.MG: 'lookupMailGroup', dns.MR: 'lookupMailRename', dns.NULL: 'lookupNull', dns.WKS: 'lookupWellKnownServices', dns.PTR: 'lookupPointer', dns.HINFO: 'lookupHostInfo', dns.MINFO: 'lookupMailboxInfo', dns.MX: 'lookupMailExchange', dns.TXT: 'lookupText', dns.SPF: 'lookupSenderPolicy', dns.RP: 'lookupResponsibility', dns.AFSDB: 'lookupAFSDatabase', dns.SRV: 'lookupService', dns.NAPTR: 'lookupNamingAuthorityPointer', dns.AXFR: 'lookupZone', dns.ALL_RECORDS: 'lookupAllRecords', } smartType = ('lookupAddress', 'lookupAuthority') class FailureHandler: def __init__(self, resolver, query, timeout, addr=None, edns=None): self.resolver = resolver self.query = query self.timeout = timeout self.addr = addr self.edns = edns def __call__(self, failure): # AuthoritativeDomainErrors should halt resolution attempts failure.trap(dns.DomainError, defer.TimeoutError, NotImplementedError) return self.resolver(self.query, self.timeout, self.addr, self.edns) class MapResolver(client.Resolver): def __init__(self, finder, a_mapping, ns_mapping, soa_mapping, servers): self.finder = finder self.a_mapping = a_mapping self.ns_mapping = ns_mapping self.soa_mapping = soa_mapping client.Resolver.__init__(self, servers=servers) def query(self, query, timeout=None, addr=None, edns=None): try: if typeToMethod[query.type] in smartType: return self.typeToMethod[query.type](str(query.name), timeout, addr, edns) else: return self.typeToMethod[query.type](str(query.name), timeout) except KeyError as e: return defer.fail(failure.Failure(NotImplementedError(str(self.__class__) + " " + str(query.type)))) def lookupAddress(self, name, timeout=None, addr=None, edns=None): def packResult(value, ttl): ret = [] add = [] for x in value: ret.append(dns.RRHeader(name, dns.A, dns.IN, ttl, dns.Record_A(x, ttl), True)) if edns is not None: if edns.rdlength > 8: add.append(dns.RRHeader('', sdns.EDNS, 4096, edns.ttl, edns.payload, True)) else: add.append(dns.RRHeader('', sdns.EDNS, 4096, 0, sdns.Record_EDNS(None, 0), True)) return [ret, (), add] wildcard = name[name.index("."):] if "." in name else None if name in self.a_mapping: ttl = self.a_mapping[name]['ttl'] result = self.finder.findIP(str(addr[0]), name) random.shuffle(result) # 返回的IP数组乱序 return packResult(result, ttl) elif wildcard is not None and wildcard in self.a_mapping: ttl = self.a_mapping[wildcard]['ttl'] result = self.finder.findIP(str(addr[0]), wildcard) random.shuffle(result) # 返回的IP数组乱序 return packResult(result, ttl) else: return self._lookup(name, dns.IN, dns.A, timeout) def lookupNameservers(self, name, timeout=None): if name in self.ns_mapping: result = self.ns_mapping[name] ttl = result['ttl'] record = re.split(r',|\s+', result['record']) def packResultNS(value): ret = [] for x in value: ret.append(dns.RRHeader(name, dns.NS, dns.IN, ttl, dns.Record_NS(x, ttl), True)) return [ret, (), ()] return packResultNS(record) else: return self._lookup(name, dns.IN, dns.NS, timeout) def lookupAuthority(self, name, timeout=None, addr=None, edns=None): if name in self.soa_mapping: result = self.soa_mapping[name] add = [] def packResultSOA(value): if edns is not None: if edns.rdlength > 8: add.append(dns.RRHeader('', dns.EDNS, 4096, edns.ttl, edns.payload, True)) else: add.append(dns.RRHeader('', dns.EDNS, 4096, 0, sdns.Record_EDNS(None, 0), True)) return [(dns.RRHeader(name, dns.SOA, dns.IN, value['ttl'], dns.Record_SOA(value['record'], value['email'], value['serial'], value['refresh'], value['retry'], value['expire'], value['ttl']), True),), (), add ] ret = packResultSOA(result) logger.info("SOA\t[domain: %s]\t[return: %s]\t[additional: %s]" % \ (name, result, add)) return ret else: return self._lookup(name, dns.IN, dns.SOA, timeout) def lookupIPV6Address(self, name, timeout=None, addr=None): return [(), (), ()] class SmartResolverChain(resolve.ResolverChain): def __init__(self, resolvers): # resolve.ResolverChain.__init__(self, resolvers) common.ResolverBase.__init__(self) self.resolvers = resolvers def _lookup(self, name, cls, type, timeout, addr=None, edns=None): q = dns.Query(name, type, cls) d = defer.fail(failure.Failure(dns.DomainError(name))) for r in self.resolvers: d = d.addErrback( FailureHandler(r.query, q, timeout, addr, edns) ) return d def _query(self, query, timeout=None, addr=None, edns=None): if typeToMethod[query.type] in smartType: return self.typeToMethod[query.type](str(query.name), timeout, addr, edns) else: return self.typeToMethod[query.type](str(query.name), timeout) def query(self, query, timeout=None, addr=None, edns=None): try: return self._query(query, timeout, addr, edns) except KeyError as e: return defer.fail(failure.Failure(NotImplementedError(str(self.__class__) + " " + str(query.type)))) def lookupAddress(self, name, timeout=None, addr=None, edns=None): return self._lookup(name, dns.IN, dns.A, timeout, addr, edns) def lookupAuthority(self, name, timeout=None, addr=None, edns=None): return self._lookup(name, dns.IN, dns.SOA, timeout, addr, edns) def lookupIPV6Address(self, name, timeout=None, addr=None, edns=None): return self._lookup(name, dns.IN, dns.AAAA, timeout, addr, edns) def lookupNameservers(self, name, timeout=None, addr=None, edns=None): return self._lookup(name, dns.IN, dns.NS, timeout, addr, edns) class SmartDNSFactory(server.DNSServerFactory): def handleQuery(self, message, protocol, address): # if len(message.additional) > 0: # print inspect.getmembers(message.additional[0] # 可以支持多个query query = message.queries[0] edns = None cliAddr = address if query.type == 43 or typeToMethod[query.type] == 'lookupAllRecords': return [(), (), ()] if isinstance(protocol, dns.DNSProtocol): cliAddr = protocol.transport.client elif typeToMethod[query.type] in smartType and \ len(message.additional) != 0 and \ message.additional[0].type == 41 and \ message.additional[0].rdlength > 8: if isinstance(message.additional[0].payload, dns.Record_A): cliAddr = (message.additional[0].payload.dottedQuad(), 0) edns = message.additional[0] return self.resolver.query(query, addr=cliAddr, edns=edns).addCallback( self.gotResolverResponse, protocol, message, address ).addErrback( self.gotResolverError, protocol, message, address ) def __init__(self, authorities=None, clients=None, verbose=0): resolvers = [] if authorities is not None: resolvers.extend(authorities) if clients is not None: resolvers.extend(clients) self.canRecurse = not not clients self.resolver = SmartResolverChain(resolvers) self.verbose = verbose self.connections = []