Add socket.setdefaulttimeout example to Resolver doctest.
[apachelog.git] / apachelog / resolve.py
1 import re as _re
2 import socket as _socket
3
4
5 class Resolver (object):
6     """A simple reverse-DNS resolver.
7
8     Maintains a class-level cache of resolved IPs to avoid repeated
9     lookups on the same IP address.
10
11     Avoid hanging if we can't resolve a name.
12
13     >>> import socket
14     >>> if hasattr(_socket, 'setdefaulttimeout'):
15     >>>     socket.setdefaulttimeout(5)  # set 5 second timeout
16
17     >>> r = Resolver()
18     >>> r.resolve('198.41.0.4')
19     'a.root-servers.net'
20     >>> r.IP
21     {'198.41.0.4': ('a.root-servers.net', [], ['198.41.0.4'])}
22
23     If you want to give shorter names to various DNS names, you can
24     add an entry to the class-level ``REGEXPS``.  The entry should use
25     your name as the key, and a list of matching regexps as the value.
26     You need to enable this enhanced resolution using the ``smart``
27     argument.
28
29     >>> r.resolve('66.249.68.33')
30     'crawl-66-249-68-33.googlebot.com'
31     >>> r = Resolver(smart=True)
32     >>> r.resolve('66.249.68.34')
33     'googlebot'
34     """
35     IP = {}
36
37     REGEXPS = {
38         'feedburner': [_re.compile('.*rate-limited-proxy-.*.google.com.*')],
39         }
40     for bot in [
41         'googlebot',
42         'yandex',
43         'baiduspider',
44         'msnbot',
45         ]:
46         REGEXPS[bot] = [_re.compile('.*{}.*'.format(bot))]
47
48     def __init__(self, smart=False):
49         self._smart = smart
50
51     def resolve(self, ip):
52         if ip not in self.IP:
53             try:
54                 self.IP[ip] = _socket.gethostbyaddr(ip)
55             except _socket.herror as e:
56                 self.IP[ip] = (ip, [], [ip])
57             except _socket.gaierror as e:
58                 self.IP[ip] = (ip, [], [ip])
59             else:
60                 if self._smart:
61                     self._smart_resolve(ip)
62         return self.IP[ip][0]
63
64     def _smart_resolve(self, ip):
65         x = self.IP[ip]
66         if self._smart:
67             for name,regexps in self.REGEXPS.items():
68                 for regexp in regexps:
69                     if regexp.match(self.IP[ip][0]):
70                         self.IP[ip] = (name, x[1], x[2])
71
72     def ips(self, name):
73         "Return a set of IP addresses used by a smart-resolved name."
74         ips = set()
75         for ip,values in self.IP.items():
76             if values[0] == name:
77                 for x in values[2]:
78                     ips.add(x)
79         return ips