dns_lookups.py

一个java写的proxy的例子

            del self.nameservers[0]
            # We keep a list of all the requests so we can cancel their
            # DNS lookup when any one of them answers
            self.requests.append(
                DnsLookupConnection(nameserver, self.hostname,
                                    self.handle_dns))
            self.outstanding_requests = self.outstanding_requests + 1
            self.requests[-1].TIMEOUT = self.outstanding_requests * 2
            
            if self.nameservers:
                # Let's create another one soon
                make_timer(1, self.issue_request)
                
    def handle_dns(self, hostname, answer):
        self.outstanding_requests = self.outstanding_requests - 1

        if not self.callback: return
        
        if not answer.isError():
            self.callback(hostname, answer)
            self.cancel()
        elif not self.outstanding_requests:
            self.issue_request()
            
class DnsLookupConnection(Connection):
    "Look up a name by contacting a single nameserver"
    # Switch from UDP to TCP after some time
    PORT = 53
    TIMEOUT = 2 # Resend the request every second
    accepts_tcp = {} # Map nameserver to 0/1, whether it accepts TCP requests
    
    def __init__(self, nameserver, hostname, callback):
        self.hostname = hostname
        self.callback = callback
        self.nameserver = nameserver
        self.retries = 0
        self.conntype = 'udp'
        Connection.__init__(self)
        try:
            self.establish_connection()
        except socket.error:
            # We couldn't even connect .. bah!
            callback(hostname, DnsResponse('error', 'could not connect to DNS server'))
            self.callback = None

    def establish_connection(self):
        if self.conntype == 'tcp':
            self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
            self.connect((self.nameserver, self.PORT))
            make_timer(30, self.handle_connect_timeout)
            # NOTE: we have to fill the buffer because otherwise we
            # won't consider this object writable, and we will never
            # call handle_connect.  This needs to be fixed somehow.
            self.send_dns_request()
        else:
            self.create_socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.connect((self.nameserver, self.PORT))
            self.send_dns_request()
        
    def __repr__(self):
        where = ''
        if self.nameserver != DnsConfig.nameservers[0]:
            where = ' @ %s' % self.nameserver
        retry = ''
        if self.retries != 0:
            retry = ' retry #%s' % self.retries
        conntype = ''
        if self.conntype == 'tcp':
            conntype = 'TCP'
        return '<%s %3s  %s%s%s>' % (color(2, 'dns-lookup'), conntype, self.hostname, retry, where)

    def cancel(self):
        if self.callback:
            if self.connected: self.close()
            self.callback = None
        
    def handle_connect(self):
        # For TCP requests only
        DnsLookupConnection.accepts_tcp[self.nameserver] = 1
    
    def handle_connect_timeout(self):
        # We're trying to perform a TCP connect
        if self.callback and not self.connected:
            DnsLookupConnection.accepts_tcp[self.nameserver] = 0
            self.callback(self.hostname, DnsResponse('error', 'timed out connecting .. %s' %
                                                     self))
            self.callback = None
            return
            
    def send_dns_request(self):
        # Issue the request and set a timeout
        if not self.callback: return # Only issue if we have someone waiting

        msg = dns.dnslib.Mpacker()
        msg.addHeader(0, 0, dns.dnsopcode.QUERY, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0)
        msg.addQuestion(self.hostname, dns.dnstype.A, dns.dnsclass.IN)
        msg = msg.getbuf()
        
        if self.conntype == 'tcp':
            self.send_buffer = dns.dnslib.pack16bit(len(msg))+msg
        else:
            self.send_buffer = msg
        make_timer(self.TIMEOUT + 0.2*self.retries, self.handle_timeout)

    def handle_timeout(self):
        # The DNS server hasn't responded to us, or we've lost the
        # packet somewhere, so let's try it again, unless the retry
        # count is too large.  Each time we retry, we increase the
        # timeout (see send_dns_request).
        if not self.callback:
            return # It's already handled, so ignore this
        if not self.connected:
            self.callback(self.hostname, DnsResponse('error', 'timed out connecting'))
            self.callback = None
            return
        self.retries = self.retries + 1
        if (self.conntype == 'udp' and
            self.accepts_tcp.get(self.nameserver, 1) and
            self.retries == 1):
            # Switch to TCP
            self.TIMEOUT = 10
            self.close()
            self.conntype = 'tcp'
            self.establish_connection()
        elif self.retries < 5:
            self.send_dns_request()
        elif self.conntype == 'udp' and self.retries < 12:
            self.send_dns_request()
        else:
            if self.callback:
                self.callback(self.hostname, DnsResponse('error', 'timed out'))
                self.callback = None
            if self.connected: self.close()

    def process_read(self):
        # Assume that the entire answer comes in one packet
        if self.conntype == 'tcp':
            if len(self.recv_buffer) < 2: return
            header = self.recv_buffer[:2]
            count = dns.dnslib.unpack16bit(header)
            if len(self.recv_buffer) < 2+count: return
            self.read(2) # header
            data = self.read(count)
            self.socket.shutdown(1)
        else:
            data = self.read(1024)
            
        msg = dns.dnslib.Munpacker(data)
        (id, qr, opcode, aa, tc, rd, ra, z, rcode,
         qdcount, ancount, nscount, arcount) = msg.getHeader()
        if tc:
            self.handle_error(socket.error,
                              (84, 'Truncated DNS packet: %s from %s for %s' %
                               (tc, self.nameserver, self.hostname)))
            return
        if rcode:
            if self.callback:
                callback, self.callback = self.callback, None
                callback(self.hostname, DnsResponse('error', 'not found .. %s' % self))
            self.close()
            return

        for i in range(qdcount):
            hostname, _, _ = msg.getQuestion()
            if hostname == self.hostname:
                # This DOES answer the question we asked
                break
        else:
            # Oops, this doesn't answer the right question.  This can
            # happen because we're using UDP, and UDP replies might end
            # up in the wrong place: open conn A, send question to A,
            # timeout, send question to A, receive answer, close our
            # object, then open a new conn B, send question to B,
            # but get the OLD answer to A as a reply.  This doesn't happen
            # with TCP but then TCP is slower.

            # Anyway, if this is the answer to a different question,
            # we ignore this read, and let the timeout take its course
            return
        
        ip_addrs = []
        for i in range(ancount):
            name, type, klass, ttl, rdlength = msg.getRRheader()
            mname = 'get%sdata' % dns.dnstype.typestr(type)
            if hasattr(msg, mname): data = getattr(msg, mname)()
            else: data = msg.getbytes(rdlength)
            if type == dns.dnstype.A:
                ip_addrs.append(data)
            if type == dns.dnstype.CNAME:
                # NOTE: should we do anything with CNAMEs?
                #message(2, 'cname record', None, None, self.hostname, '=', repr(data))
                pass
            
        # Ignore (nscount) authority records
        # Ignore (arcount) additional records
        if self.callback:
            callback, self.callback = self.callback, None
            if ip_addrs: callback(self.hostname, DnsResponse('found', ip_addrs))
            else:        callback(self.hostname, DnsResponse('error', 'not found'))
        self.close()

    def handle_error(self, type, value, traceback=None):
        Connection.handle_error(self, type, value, traceback)
        if self.callback:
            callback, self.callback = self.callback, None
            callback(self.hostname, DnsResponse('error', 'failed lookup .. %s' % self))
        
    def handle_close(self):
        # If we ever get here, we want to make sure we notify the
        # callbacks so that they don't get stuck
        Connection.handle_close(self)
        if self.callback:
            callback, self.callback = self.callback, None
            callback(self.hostname, DnsResponse('error', 'closed with no answer .. %s' % self))

def init_dns_resolver():
    "Set up the DnsLookupConnection class with /etc/resolv.conf information"
    for line in open('/etc/resolv.conf', 'r').readlines():
        m = re.match(r'^search\s+\.?(.*)$', line)
        if m:
            for domain in split(m.group(1)):
                DnsConfig.search_domains.append('.'+lower(domain))
        m = re.match(r'^nameserver\s+(\S+)\s*$', line)
        if m: DnsConfig.nameservers.append(m.group(1))
    if not DnsConfig.search_domains:
        DnsConfig.search_domains.append('')
    if not DnsConfig.nameservers:
        print 'Warning: no nameservers found'
        DnsConfig.nameservers.append('127.0.0.1')


init_dns_resolver()
dnscache = DnsCache()