client/cros/backchannel.py

# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging, os

from autotest_lib.client.bin import utils
from autotest_lib.client.common_lib import error

# Flag file used to tell backchannel script it's okay to run.
BACKCHANNEL_FILE = '/mnt/stateful_partition/etc/enable_backchannel_network'


class Backchannel(object):
    """Wrap backchannel in a context manager so it can be used with with.

    Example usage:
         with backchannel.Backchannel():
                block
    The backchannel will be torn down whether or not 'block' throws.
    """

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self.interface = None

    def __enter__(self):
        self.setup(*self.args, **self.kwargs)
        return self

    def __exit__(self, exception, value, traceback):
        self.teardown()
        return False

    def setup(self, create_ssh_routes=True):
        """Enables the backchannel interface.

        Args:
          create_ssh_routes: If true set up routes so that all
          existing SSH sessions will remain open.

        Returns:
            True if the backchannel is already set up, or was
            set up by this call, otherwise False.

        """

        # If the backchannel interface is already up there's nothing
        # for us to do.
        if is_network_iface_running('eth_test'):
            return True

        # Retrieve the gateway for the default route.
        try:
            line = utils.system_output(
                "route -n | awk '/^0.0.0.0/ { print $2, $8 }'").split('\n')[0]
            gateway, self.interface = line.strip().split(' ')

            # Retrieve list of open ssh sessions so we can reopen
            # routes afterward.
            if create_ssh_routes:
                out = utils.system_output(
                    "netstat -tanp | grep :22 | "
                    "grep ESTABLISHED | awk '{print $5}'")

                # Extract IP from IP:PORT listing. Uses set to remove
                # duplicates.
                open_ssh = list(set(item.strip().split(':')[0] for item in
                                    out.split('\n') if item.strip()))

            backchannel('setup %s' % self.interface)

            # Create routes so existing SSH sessions will stay open.
            if create_ssh_routes:
                for ip in open_ssh:
                    # Add route using the pre-backchannel gateway.
                    backchannel('reach %s %s' % (ip, gateway))

            # Make sure we have a route to the gateway before continuing.
            logging.info('Waiting for route to gateway %s' % gateway)
            utils.poll_for_condition(
                lambda: is_route_ready(gateway),
                exception=utils.TimeoutError('Timed out waiting for route'),
                timeout=30)
        except Exception, e:
            logging.error(e)
            return False
        finally:
            # Remove backchannel file flag so system reverts to normal
            # on reboot.
            if os.path.isfile(BACKCHANNEL_FILE):
                os.remove(BACKCHANNEL_FILE)

        return True

    def teardown(self):
        if self.interface:
            backchannel('teardown %s' % self.interface)


def backchannel(args):
    utils.system('/usr/local/lib/flimflam/test/backchannel %s' % args)


def is_network_iface_running(name):
    try:
        out = utils.system_output('ifconfig %s' % name)
    except error.CmdError, e:
        logging.info(e)
        return False

    return out.find('RUNNING') >= 0


def is_route_ready(dest):
    try:
        out = utils.system_output('ping -c 1 %s' % dest)
    except error.CmdError, e:
        logging.error(e)
        return False

    return True