1# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import common, constants, logging, os, socket, stat, sys, threading, time
6
7from autotest_lib.client.bin import utils
8from autotest_lib.client.common_lib import error
9
10class LocalDns(object):
11    """A wrapper around miniFakeDns that runs the server in a separate thread
12    and redirects all DNS queries to it.
13    """
14    # This is a symlink.  We look up the real path at runtime by following it.
15    _resolv_bak_file = 'resolv.conf.bak'
16
17    def __init__(self, fake_ip="127.0.0.1", local_port=53):
18        import miniFakeDns  # So we don't need to install it in the chroot.
19        self._dns = miniFakeDns.DNSServer(fake_ip=fake_ip, port=local_port)
20        self._stopper = threading.Event()
21        self._thread = threading.Thread(target=self._dns.run,
22                                        args=(self._stopper,))
23
24    def __get_host_by_name(self, hostname):
25        """Resolve the dotted-quad IPv4 address of |hostname|
26
27        This used to use suave python code, like this:
28            hosts = socket.getaddrinfo(hostname, 80, socket.AF_INET)
29            (fam, socktype, proto, canonname, (host, port)) = hosts[0]
30            return host
31
32        But that hangs sometimes, and we don't understand why.  So, use
33        a subprocess with a timeout.
34        """
35        try:
36            host = utils.system_output('%s -c "import socket; '
37                                       'print socket.gethostbyname(\'%s\')"' % (
38                                       sys.executable, hostname),
39                                       ignore_status=True, timeout=2)
40        except Exception as e:
41            logging.warning(e)
42            return None
43        return host or None
44
45    def __attempt_resolve(self, hostname, ip, expected=True):
46        logging.debug('Attempting to resolve %s to %s' % (hostname, ip))
47        host = self.__get_host_by_name(hostname)
48        logging.debug('Resolve attempt for %s got %s' % (hostname, host))
49        return host and (host == ip) == expected
50
51    def run(self):
52        """Start the mock DNS server and redirect all queries to it."""
53        self._thread.start()
54        # Redirect all DNS queries to the mock DNS server.
55        try:
56            # Follow resolv.conf symlink.
57            resolv = os.path.realpath(constants.RESOLV_CONF_FILE)
58            # Grab path to the real file, do following work in that directory.
59            resolv_dir = os.path.dirname(resolv)
60            resolv_bak = os.path.join(resolv_dir, self._resolv_bak_file)
61            resolv_contents = 'nameserver 127.0.0.1'
62            # Test to make sure the current resolv.conf isn't already our
63            # specially modified version.  If this is the case, we have
64            # probably been interrupted while in the middle of this test
65            # in a previous run.  The last thing we want to do at this point
66            # is to overwrite a legitimate backup.
67            if (utils.read_one_line(resolv) == resolv_contents and
68                os.path.exists(resolv_bak)):
69                logging.error('Current resolv.conf is setup for our local '
70                              'server, and a backup already exists!  '
71                              'Skipping the backup step.')
72            else:
73                # Back up the current resolv.conf.
74                os.rename(resolv, resolv_bak)
75            # To stop flimflam from editing resolv.conf while we're working
76            # with it, we want to make the directory -r-xr-xr-x.  Open an
77            # fd to the file first, so that we'll retain the ability to
78            # alter it.
79            resolv_fd = open(resolv, 'w')
80            self._resolv_dir_mode = os.stat(resolv_dir).st_mode
81            os.chmod(resolv_dir, (stat.S_IRUSR | stat.S_IXUSR |
82                                  stat.S_IRGRP | stat.S_IXGRP |
83                                  stat.S_IROTH | stat.S_IXOTH))
84            resolv_fd.write(resolv_contents)
85            resolv_fd.close()
86            assert utils.read_one_line(resolv) == resolv_contents
87        except Exception as e:
88            logging.error(str(e))
89            raise e
90
91        utils.poll_for_condition(
92            lambda: self.__attempt_resolve('www.google.com.', '127.0.0.1'),
93            utils.TimeoutError('Timed out waiting for DNS changes.'),
94            timeout=10)
95
96    def stop(self):
97        """Restore the backed-up DNS settings and stop the mock DNS server."""
98        try:
99            # Follow resolv.conf symlink.
100            resolv = os.path.realpath(constants.RESOLV_CONF_FILE)
101            # Grab path to the real file, do following work in that directory.
102            resolv_dir = os.path.dirname(resolv)
103            resolv_bak = os.path.join(resolv_dir, self._resolv_bak_file)
104            os.chmod(resolv_dir, self._resolv_dir_mode)
105            if os.path.exists(resolv_bak):
106                os.rename(resolv_bak, resolv)
107            else:
108                # This probably means shill restarted during the execution
109                # of our test, and has cleaned up the .bak file we created.
110                raise error.TestError('Backup file %s no longer exists!  '
111                                      'Connection manager probably crashed '
112                                      'during the test run.' %
113                                      resolv_bak)
114
115            utils.poll_for_condition(
116                lambda: self.__attempt_resolve('www.google.com.',
117                                               '127.0.0.1',
118                                               expected=False),
119                utils.TimeoutError('Timed out waiting to revert DNS.  '
120                                   'resolv.conf contents are: ' +
121                                   utils.read_one_line(resolv)),
122                timeout=10)
123        finally:
124            # Stop the DNS server.
125            self._stopper.set()
126            self._thread.join()
127