Permalink
Switch branches/tags
Find file
Fetching contributors…
Cannot retrieve contributors at this time
229 lines (199 sloc) 7.31 KB
#!/usr/bin/env python
# Copyright 2010 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import daemonserver
import errno
import logging
import socket
import SocketServer
import threading
import third_party
import dns.flags
import dns.message
import dns.rcode
import dns.resolver
import dns.rdatatype
import ipaddr
class DnsProxyException(Exception):
pass
class RealDnsLookup(object):
def __init__(self, name_servers):
if '127.0.0.1' in name_servers:
raise DnsProxyException(
'Invalid nameserver: 127.0.0.1 (causes an infinte loop)')
self.resolver = dns.resolver.get_default_resolver()
self.resolver.nameservers = name_servers
self.dns_cache_lock = threading.Lock()
self.dns_cache = {}
def __call__(self, hostname, rdtype=dns.rdatatype.A):
"""Return real IP for a host.
Args:
host: a hostname ending with a period (e.g. "www.google.com.")
rdtype: the query type (1 for 'A', 28 for 'AAAA')
Returns:
the IP address as a string (e.g. "192.168.25.2")
"""
self.dns_cache_lock.acquire()
ip = self.dns_cache.get(hostname)
self.dns_cache_lock.release()
if ip:
return ip
try:
answers = self.resolver.query(hostname, rdtype)
except dns.resolver.NXDOMAIN:
return None
except (dns.resolver.NoAnswer, dns.resolver.Timeout) as ex:
logging.debug('_real_dns_lookup(%s) -> None (%s)',
hostname, ex.__class__.__name__)
return None
if answers:
ip = str(answers[0])
self.dns_cache_lock.acquire()
self.dns_cache[hostname] = ip
self.dns_cache_lock.release()
return ip
def ClearCache(self):
"""Clearn the dns cache."""
self.dns_cache_lock.acquire()
self.dns_cache.clear()
self.dns_cache_lock.release()
class PrivateIpDnsLookup(object):
"""Resolve private hosts to their real IPs and others to the Web proxy IP.
Hosts in the given http_archive will resolve to the Web proxy IP without
checking the real IP.
This only supports IPv4 lookups.
"""
def __init__(self, web_proxy_ip, real_dns_lookup, http_archive):
"""Initialize PrivateIpDnsLookup.
Args:
web_proxy_ip: the IP address returned by __call__ for non-private hosts.
real_dns_lookup: a function that resolves a host to an IP.
http_archive: an instance of a HttpArchive
Hosts is in the archive will always resolve to the web_proxy_ip
"""
self.web_proxy_ip = web_proxy_ip
self.real_dns_lookup = real_dns_lookup
self.http_archive = http_archive
self.InitializeArchiveHosts()
def __call__(self, host):
"""Return real IPv4 for private hosts and Web proxy IP otherwise.
Args:
host: a hostname ending with a period (e.g. "www.google.com.")
Returns:
IP address as a string or None (if lookup fails)
"""
ip = self.web_proxy_ip
if host not in self.archive_hosts:
real_ip = self.real_dns_lookup(host)
if real_ip:
if ipaddr.IPAddress(real_ip).is_private:
ip = real_ip
else:
ip = None
return ip
def InitializeArchiveHosts(self):
"""Recompute the archive_hosts from the http_archive."""
self.archive_hosts = set('%s.' % req.host for req in self.http_archive)
class UdpDnsHandler(SocketServer.DatagramRequestHandler):
"""Resolve DNS queries to localhost.
Possible alternative implementation:
http://howl.play-bow.org/pipermail/dnspython-users/2010-February/000119.html
"""
STANDARD_QUERY_OPERATION_CODE = 0
def handle(self):
"""Handle a DNS query.
IPv6 requests (with rdtype AAAA) receive mismatched IPv4 responses
(with rdtype A). To properly support IPv6, the http proxy would
need both types of addresses. By default, Windows XP does not
support IPv6.
"""
self.data = self.rfile.read()
self.transaction_id = self.data[0]
self.flags = self.data[1]
self.qa_counts = self.data[4:6]
self.domain = ''
operation_code = (ord(self.data[2]) >> 3) & 15
if operation_code == self.STANDARD_QUERY_OPERATION_CODE:
self.wire_domain = self.data[12:]
self.domain = self._domain(self.wire_domain)
else:
logging.debug("DNS request with non-zero operation code: %s",
operation_code)
ip = self.server.dns_lookup(self.domain)
if ip is None:
logging.debug('dnsproxy: %s -> NXDOMAIN', self.domain)
response = self.get_dns_no_such_name_response()
else:
if ip == self.server.server_address[0]:
logging.debug('dnsproxy: %s -> %s (replay web proxy)', self.domain, ip)
else:
logging.debug('dnsproxy: %s -> %s', self.domain, ip)
response = self.get_dns_response(ip)
self.wfile.write(response)
@classmethod
def _domain(cls, wire_domain):
domain = ''
index = 0
length = ord(wire_domain[index])
while length:
domain += wire_domain[index + 1:index + length + 1] + '.'
index += length + 1
length = ord(wire_domain[index])
return domain
def get_dns_response(self, ip):
packet = ''
if self.domain:
packet = (
self.transaction_id +
self.flags +
'\x81\x80' + # standard query response, no error
self.qa_counts * 2 + '\x00\x00\x00\x00' + # Q&A counts
self.wire_domain +
'\xc0\x0c' # pointer to domain name
'\x00\x01' # resource record type ("A" host address)
'\x00\x01' # class of the data
'\x00\x00\x00\x3c' # ttl (seconds)
'\x00\x04' + # resource data length (4 bytes for ip)
socket.inet_aton(ip)
)
return packet
def get_dns_no_such_name_response(self):
query_message = dns.message.from_wire(self.data)
response_message = dns.message.make_response(query_message)
response_message.flags |= dns.flags.AA | dns.flags.RA
response_message.set_rcode(dns.rcode.NXDOMAIN)
return response_message.to_wire()
class DnsProxyServer(SocketServer.ThreadingUDPServer,
daemonserver.DaemonServer):
def __init__(self, dns_lookup=None, host='', port=53):
"""Initialize DnsProxyServer.
Args:
dns_lookup: a function that resolves a host to an IP address.
host: a host string (name or IP) to bind the dns proxy and to which
DNS requests will be resolved.
port: an integer port on which to bind the proxy.
"""
try:
SocketServer.ThreadingUDPServer.__init__(
self, (host, port), UdpDnsHandler)
except socket.error, (error_number, msg):
if error_number == errno.EACCES:
raise DnsProxyException(
'Unable to bind DNS server on (%s:%s)' % (host, port))
raise
self.dns_lookup = dns_lookup or (lambda host: self.server_address[0])
logging.info('Started DNS server on %s...', self.server_address)
def cleanup(self):
self.shutdown()
logging.info('Shutdown DNS server')