src/third_party/webpagereplay/dnsproxy.py

   1 #!/usr/bin/env python
   2 # Copyright 2010 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #      http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 import daemonserver
  17 import errno
  18 import logging
  19 import socket
  20 import SocketServer
  21 import threading
  22 import time
  23
  24 import third_party
  25 import dns.flags
  26 import dns.message
  27 import dns.rcode
  28 import dns.resolver
  29 import dns.rdatatype
  30 import ipaddr
  31
  32
  33 class DnsProxyException(Exception):
  34   pass
  35
  36
  37 class RealDnsLookup(object):
  38   def __init__(self, name_servers):
  39     if '127.0.0.1' in name_servers:
  40       raise DnsProxyException(
  41           'Invalid nameserver: 127.0.0.1 (causes an infinte loop)')
  42     self.resolver = dns.resolver.get_default_resolver()
  43     self.resolver.nameservers = name_servers
  44     self.dns_cache_lock = threading.Lock()
  45     self.dns_cache = {}
  46
  47   def _IsIPAddress(self, hostname):
  48     try:
  49       socket.inet_aton(hostname)
  50       return True
  51     except socket.error:
  52       return False
  53
  54   def __call__(self, hostname, rdtype=dns.rdatatype.A):
  55     """Return real IP for a host.
  56
  57     Args:
  58       host: a hostname ending with a period (e.g. "www.google.com.")
  59       rdtype: the query type (1 for 'A', 28 for 'AAAA')
  60     Returns:
  61       the IP address as a string (e.g. "192.168.25.2")
  62     """
  63     if self._IsIPAddress(hostname):
  64       return hostname
  65     self.dns_cache_lock.acquire()
  66     ip = self.dns_cache.get(hostname)
  67     self.dns_cache_lock.release()
  68     if ip:
  69       return ip
  70     try:
  71       answers = self.resolver.query(hostname, rdtype)
  72     except dns.resolver.NXDOMAIN:
  73       return None
  74     except dns.resolver.NoNameservers:
  75       logging.debug('_real_dns_lookup(%s) -> No nameserver.',
  76                     hostname)
  77       return None
  78     except (dns.resolver.NoAnswer, dns.resolver.Timeout) as ex:
  79       logging.debug('_real_dns_lookup(%s) -> None (%s)',
  80                     hostname, ex.__class__.__name__)
  81       return None
  82     if answers:
  83       ip = str(answers[0])
  84     self.dns_cache_lock.acquire()
  85     self.dns_cache[hostname] = ip
  86     self.dns_cache_lock.release()
  87     return ip
  88
  89   def ClearCache(self):
  90     """Clear the dns cache."""
  91     self.dns_cache_lock.acquire()
  92     self.dns_cache.clear()
  93     self.dns_cache_lock.release()
  94
  95
  96 class ReplayDnsLookup(object):
  97   """Resolve DNS requests to replay host."""
  98   def __init__(self, replay_ip, filters=None):
  99     self.replay_ip = replay_ip
 100     self.filters = filters or []
 101
 102   def __call__(self, hostname):
 103     ip = self.replay_ip
 104     for f in self.filters:
 105       ip = f(hostname, default_ip=ip)
 106     return ip
 107
 108
 109 class PrivateIpFilter(object):
 110   """Resolve private hosts to their real IPs and others to the Web proxy IP.
 111
 112   Hosts in the given http_archive will resolve to the Web proxy IP without
 113   checking the real IP.
 114
 115   This only supports IPv4 lookups.
 116   """
 117   def __init__(self, real_dns_lookup, http_archive):
 118     """Initialize PrivateIpDnsLookup.
 119
 120     Args:
 121       real_dns_lookup: a function that resolves a host to an IP.
 122       http_archive: an instance of a HttpArchive
 123         Hosts is in the archive will always resolve to the web_proxy_ip
 124     """
 125     self.real_dns_lookup = real_dns_lookup
 126     self.http_archive = http_archive
 127     self.InitializeArchiveHosts()
 128
 129   def __call__(self, host, default_ip):
 130     """Return real IPv4 for private hosts and Web proxy IP otherwise.
 131
 132     Args:
 133       host: a hostname ending with a period (e.g. "www.google.com.")
 134     Returns:
 135       IP address as a string or None (if lookup fails)
 136     """
 137     ip = default_ip
 138     if host not in self.archive_hosts:
 139       real_ip = self.real_dns_lookup(host)
 140       if real_ip:
 141         if ipaddr.IPAddress(real_ip).is_private:
 142           ip = real_ip
 143       else:
 144         ip = None
 145     return ip
 146
 147   def InitializeArchiveHosts(self):
 148     """Recompute the archive_hosts from the http_archive."""
 149     self.archive_hosts = set('%s.' % req.host for req in self.http_archive)
 150
 151
 152 class DelayFilter(object):
 153   """Add a delay to replayed lookups."""
 154
 155   def __init__(self, is_record_mode, delay_ms):
 156     self.is_record_mode = is_record_mode
 157     self.delay_ms = int(delay_ms)
 158
 159   def __call__(self, host, default_ip):
 160     if not self.is_record_mode:
 161       time.sleep(self.delay_ms * 1000.0)
 162     return default_ip
 163
 164   def SetRecordMode(self):
 165     self.is_record_mode = True
 166
 167   def SetReplayMode(self):
 168     self.is_record_mode = False
 169
 170
 171 class UdpDnsHandler(SocketServer.DatagramRequestHandler):
 172   """Resolve DNS queries to localhost.
 173
 174   Possible alternative implementation:
 175   http://howl.play-bow.org/pipermail/dnspython-users/2010-February/000119.html
 176   """
 177
 178   STANDARD_QUERY_OPERATION_CODE = 0
 179
 180   def handle(self):
 181     """Handle a DNS query.
 182
 183     IPv6 requests (with rdtype AAAA) receive mismatched IPv4 responses
 184     (with rdtype A). To properly support IPv6, the http proxy would
 185     need both types of addresses. By default, Windows XP does not
 186     support IPv6.
 187     """
 188     self.data = self.rfile.read()
 189     self.transaction_id = self.data[0]
 190     self.flags = self.data[1]
 191     self.qa_counts = self.data[4:6]
 192     self.domain = ''
 193     operation_code = (ord(self.data[2]) >> 3) & 15
 194     if operation_code == self.STANDARD_QUERY_OPERATION_CODE:
 195       self.wire_domain = self.data[12:]
 196       self.domain = self._domain(self.wire_domain)
 197     else:
 198       logging.debug("DNS request with non-zero operation code: %s",
 199                     operation_code)
 200     ip = self.server.dns_lookup(self.domain)
 201     if ip is None:
 202       logging.debug('dnsproxy: %s -> NXDOMAIN', self.domain)
 203       response = self.get_dns_no_such_name_response()
 204     else:
 205       if ip == self.server.server_address[0]:
 206         logging.debug('dnsproxy: %s -> %s (replay web proxy)', self.domain, ip)
 207       else:
 208         logging.debug('dnsproxy: %s -> %s', self.domain, ip)
 209       response = self.get_dns_response(ip)
 210     self.wfile.write(response)
 211
 212   @classmethod
 213   def _domain(cls, wire_domain):
 214     domain = ''
 215     index = 0
 216     length = ord(wire_domain[index])
 217     while length:
 218       domain += wire_domain[index + 1:index + length + 1] + '.'
 219       index += length + 1
 220       length = ord(wire_domain[index])
 221     return domain
 222
 223   def get_dns_response(self, ip):
 224     packet = ''
 225     if self.domain:
 226       packet = (
 227           self.transaction_id +
 228           self.flags +
 229           '\x81\x80' +        # standard query response, no error
 230           self.qa_counts * 2 + '\x00\x00\x00\x00' +  # Q&A counts
 231           self.wire_domain +
 232           '\xc0\x0c'          # pointer to domain name
 233           '\x00\x01'          # resource record type ("A" host address)
 234           '\x00\x01'          # class of the data
 235           '\x00\x00\x00\x3c'  # ttl (seconds)
 236           '\x00\x04' +        # resource data length (4 bytes for ip)
 237           socket.inet_aton(ip)
 238           )
 239     return packet
 240
 241   def get_dns_no_such_name_response(self):
 242     query_message = dns.message.from_wire(self.data)
 243     response_message = dns.message.make_response(query_message)
 244     response_message.flags |= dns.flags.AA | dns.flags.RA
 245     response_message.set_rcode(dns.rcode.NXDOMAIN)
 246     return response_message.to_wire()
 247
 248
 249 class DnsProxyServer(SocketServer.ThreadingUDPServer,
 250                      daemonserver.DaemonServer):
 251   # Increase the request queue size. The default value, 5, is set in
 252   # SocketServer.TCPServer (the parent of BaseHTTPServer.HTTPServer).
 253   # Since we're intercepting many domains through this single server,
 254   # it is quite possible to get more than 5 concurrent requests.
 255   request_queue_size = 128
 256
 257   # Don't prevent python from exiting when there is thread activity.
 258   daemon_threads = True
 259
 260   def __init__(self, host='', port=53, dns_lookup=None):
 261     """Initialize DnsProxyServer.
 262
 263     Args:
 264       host: a host string (name or IP) to bind the dns proxy and to which
 265         DNS requests will be resolved.
 266       port: an integer port on which to bind the proxy.
 267       dns_lookup: a list of filters to apply to lookup.
 268     """
 269     try:
 270       SocketServer.ThreadingUDPServer.__init__(
 271           self, (host, port), UdpDnsHandler)
 272     except socket.error, (error_number, msg):
 273       if error_number == errno.EACCES:
 274         raise DnsProxyException(
 275             'Unable to bind DNS server on (%s:%s)' % (host, port))
 276       raise
 277     self.dns_lookup = dns_lookup or (lambda host: self.server_address[0])
 278     self.server_port = self.server_address[1]
 279     logging.warning('DNS server started on %s:%d', self.server_address[0],
 280                                                    self.server_address[1])
 281
 282   def cleanup(self):
 283     try:
 284       self.shutdown()
 285     except KeyboardInterrupt, e:
 286       pass
 287     logging.info('Stopped DNS server')