dns_lookups.py

一个java写的proxy的例子

from proxy4_base import *
from connection import *

# For a high level overview of DNS, see
# http://www.rad.com/networks/1998/dns/main.html

import whrandom

import dns.dnslib # Guido's DNS libraries put into a package
import dns.dnsclass 
import dns.dnsopcode
import dns.dnstype

def background_lookup(hostname, callback):
    "Return immediately, but call callback with a DnsResponse object later"
    # Hostnames are case insensitive, so canonicalize for lookup purposes
    DnsExpandHostname(lower(hostname), callback)
    
class DnsResponse:
    # A DNS answer can be:
    #    ('found', [ipaddrs])
    #    ('error', why-str)
    #    ('redirect', new-hostname)
    def __init__(self, kind, data):
        self.kind = kind
        self.data = data
        print 'DnsResponse(%s, %s)' % (`kind`, `data`)
        
    def isError(self):
        return self.kind == 'error'

    def isFound(self):
        return self.kind == 'found'

    def isRedirect(self):
        return self.kind == 'redirect'
    
class DnsConfig:
    nameservers = []
    search_domains = []
    search_patterns = ('www.%s.com', ) # 'www.%s.net', 'www.%s.org')
    
class DnsExpandHostname:
    "Try looking up a hostname and its expansions"
    # This routine calls DnsCache to do the individual lookups
    def __init__(self, hostname, callback):
        self.hostname = hostname
        self.callback = callback
        self.queries = [hostname] # all queries to be made
        self.answers = {} # Map hostname to DNS answer
        self.delay = 0.2 # How long do we wait before trying another expansion?
        if not dnscache.well_known_hosts.has_key(hostname):
            for domain in DnsConfig.search_domains:
                self.queries.append(hostname + domain)
            if find(hostname, '.') < 0:
                # If there's no dot, we should try expanding patterns
                for pattern in DnsConfig.search_patterns:
                    self.queries.append(pattern % hostname)
            else:
                # But if there is a dot, let's increase the delay
                # because it's very likely that none of the
                # search_domains matter.
                self.delay = 3
        if re.search(r'\s', hostname):
            # If there's whitespace, it's almost certainly a copy/paste error,
            # so also try the same thing with whitespace removed
            self.queries.append(re.sub(r'\s+', '', hostname))
        self.requests = self.queries[1:] # queries we haven't yet made

        # Issue the primary request
        make_timer(0, lambda h=hostname, s=self:
                   dnscache.lookup(h, s.handle_dns))
        # and then start another request as well
        if self.delay < 1: # (it's likely to be needed)
            make_timer(self.delay, self.handle_issue_request)

    def handle_issue_request(self):
        # Issue one DNS request, and set up a timer to issue another
        if self.requests and self.callback:
            request = self.requests[0]
            del self.requests[0]
            make_timer(0, lambda r=request, s=self:
                       dnscache.lookup(r, s.handle_dns))

            # NOTE: Yes, it's possible that several DNS lookups are
            # being executed at once.  To avoid that, we could check
            # if there's already a timer for this object ..
            if self.requests: make_timer(self.delay, self.handle_issue_request)
        
    def handle_dns(self, hostname, answer):
        if not self.callback: return # Already handled this query

        self.answers[hostname] = answer
        while self.queries and self.answers.has_key(self.queries[0]):
            current_query = self.queries[0]
            del self.queries[0]
            answer = self.answers[current_query]
            if not answer.isError():
                callback, self.callback = self.callback, None
                if self.hostname != current_query:
                    callback(self.hostname, DnsResponse('redirect', current_query))
                else:
                    callback(self.hostname, answer)
                break

        if self.callback and not self.queries:
            # Someone's still waiting for an answer, and we 
            # are expecting no more answers
            callback, self.callback = self.callback, None
            callback(self.hostname, DnsResponse('error', 'host not expanded'))

        # Since one DNS request is satisfied, issue another
        self.handle_issue_request()
            
class DnsCache:
    """Provides a lookup function that will either get a value from the cache or
    initiate a DNS lookup, fill the cache, and return that value"""
    # lookup() can create zero or one DnsLookupHostname objects
    
    ValidCacheEntryExpires = 1800
    InvalidCacheEntryExpires = 30
    
    def __init__(self):
        self.cache = {} # hostname to DNS answer
        self.expires = {} # hostname to cache expiry time
        self.pending = {} # hostname to [callbacks]
        self.well_known_hosts = {} # hostname to 1, if it's from /etc/hosts
        self.read_localhosts()

    def read_localhosts(self):
        "Fill DnsCache with /etc/hosts information"
        for line in open('/etc/hosts', 'r').readlines():
            if find(line, '#') >= 0: line = line[:find(line, '#')] # Comments
            fields = split(line)
            if len(fields) > 0:
                # The first one is the IP address, and then the rest are names
                # These hosts don't expire from our cache
                for name in fields[1:]:
                    name = lower(name)
                    self.well_known_hosts[name] = 1
                    self.cache[name] = DnsResponse('found', [fields[0]])
                    self.expires[name] = sys.maxint
        
    def lookup(self, hostname, callback):
        if re.match(r'^[.\d]+$', hostname):
            # It's an IP address, so let's just return the same thing
            # NOTE: this will have to be changed for IPv6
            callback(hostname, DnsResponse('found', [hostname]))
            return

        if hostname[-1:] == '.':
            # We should just remove the trailing '.'
            return DnsResponse('redirect', hostname[:-1])

        if len(hostname) > 100:
            # It's too long .. assume it's an error
            callback(hostname, DnsResponse('error', 'hostname %s too long' % hostname))
            return
        
        if self.cache.has_key(hostname):
            if time.time() < self.expires[hostname]:
                # It hasn't expired, so return this answer
                callback(hostname, self.cache[hostname])
                return
            elif not self.cache[hostname].isError():
                # It has expired, but we can use the old value for now
                callback(hostname, self.cache[hostname])
                # We *don't* return; instead, we trigger a new cache
                # fill and use a dummy callback
                callback = lambda h,a:None

        if self.pending.has_key(hostname):
            # Add this guy to the list of interested parties
            self.pending[hostname].append(callback)
            return
        else:
            # Create a new lookup object
            self.pending[hostname] = [callback]
            DnsLookupHostname(hostname, self.handle_dns)
            
    def handle_dns(self, hostname, answer):
        assert self.pending.has_key(hostname)
        callbacks = self.pending[hostname]
        del self.pending[hostname]

        assert (not answer.isFound() or len(answer.data) > 0), \
                'Received empty DNS lookup .. should be error? %s' % (answer,)
        
        self.cache[hostname] = answer
        if not answer.isError():
            self.expires[hostname] = time.time()+self.ValidCacheEntryExpires
        else:
            self.expires[hostname] = time.time()+self.InvalidCacheEntryExpires
            
        for c in callbacks:
            c(hostname, answer)

class DnsLookupHostname:
    "Perform DNS lookup on many nameservers"
    # Use a DnsLookupConnection per nameserver
    
    # We start working with one nameserver per second, as long as we
    # haven't gotten any responses.  For each successive nameserver we
    # set the timeout higher, so that the first nameserver has to try
    # harder.
    def __init__(self, hostname, callback):
        self.hostname = hostname
        self.callback = callback
        self.nameservers = DnsConfig.nameservers[:]
        self.requests = []
        self.outstanding_requests = 0
        
        self.issue_request()

    def cancel(self):
        if self.callback:
            self.callback = None
            # Now let's go through and tell all the lookup operations
            # that there's no need to contact us
            for r in self.requests:
                if r.callback == self.handle_dns:
                    r.cancel()
                    self.outstanding_requests = self.outstanding_requests - 1
                assert r.callback is None
            assert self.outstanding_requests == 0
            
    def issue_request(self):
        if not self.callback: return
        
        if not self.nameservers and not self.outstanding_requests:
            self.callback(self.hostname, DnsResponse('error', 'no nameserver found host'))
            self.callback = None
        elif self.nameservers:
            nameserver = self.nameservers[0]