path: root/chrome/test/pyautolib/chromeos_network.py

#!/usr/bin/python

# Copyright (c) 2011 The Chromium 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 copy
import dbus
import logging
import os
import time

from chromeos.power_strip import PowerStrip
import pyauto
import pyauto_errors

class WifiPowerStrip(PowerStrip):
  """Manages the power state of wifi routers connected to a power strip.

  This class provides additional functionality over PowerStrip by providing
  a timeout feature for wifi routers connected to the strip.  This is to prevent
  repeated on/off calls to the same router which may put the router in an
  undesired state.
  """

  def __init__ (self, host, routers):
    """Initializes a WifiPowerStrip object.

    Args:
      host: IP of the switch that the routers are attached to.
      routers: Dictionary of wifi routers in the following format:
               {
                 '< router name >': {
                   'strip_id' : '.aX'  # where X is the port number
                   < additional fields may be added here for each router >
                 }
               }
    """
    self._router_dict = routers

    # Each router will have a timestamp associated to it regarding whether
    # or not an action can be performed on it yet.   This is to prevent
    # the spamming of power on/off calls on a particular router.
    # The WifiPowerStrip_UsableTime field specifies the earliest time
    # after which the router may be used.  We will initialize it to now
    # since they should all be usable at init.
    for router_info in self._router_dict.values():
      router_info['WifiPowerStrip_UsableTime'] = time.time()

    # _routers_used keeps track of which routers were used during the lifetime
    # of the WifiPowerStrip instance.  Adding used routers occurs when
    # a wifi router has been turned on.  Otherwise, it get clears upon
    # the TurnOffUsedRouters call.
    self._routers_used = set()
    PowerStrip.__init__(self, host)

  def GetRouterConfig(self, router_name):
    """Returns the configuration for the specified router.

    Args:
      router_name: A string specifying the router.

    Returns:
      The config dictionary for the given router if the router is defined.
      None otherwise.
    """
    return copy.deepcopy(self._router_dict.get(router_name))

  def RouterPower(self, router_name, power_state, pause_after=5):
    """Executes PowerStrip commands.

    Args:
      router_name: The name of the router to perform the action on.
      power_state: A boolean value where True represents turning the router on
                   and False represents turning the router off.
      pause_after: Specified in seconds, and specifies the time to sleep
                   after a command is run.  This is to prevent spamming of
                   power on/off of the same router which has put the router
                   in an undesirable state.

    Raises:
      Exception if router_name is not a valid router.
    """
    router = self.GetRouterConfig(router_name)
    if not router: raise Exception('Invalid router name \'%s\'.' % router_name)

    # Hidden routers will always be on.  Don't allow controlling of the power
    # for these networks.
    if router.get('hidden'):
      return

    sleep_time = router['WifiPowerStrip_UsableTime'] - time.time()
    if sleep_time > 0:
      time.sleep(sleep_time)

    if power_state:
      self._routers_used |= set([router_name])
      logging.debug('Turning on router %s:%s.' %
                    (router['strip_id'], router_name))
      self.PowerOn(router['strip_id'])
    else:
      logging.debug('Turning off router %s:%s.' %
                    (router['strip_id'], router_name))
      self.PowerOff(router['strip_id'])

    # Set the Usable time of the particular router to pause_after
    # seconds after the current time.
    router['WifiPowerStrip_UsableTime'] = time.time() + pause_after

  def TurnOffAllRouters(self):
    """Turns off all the routers."""
    for router in self._router_dict:
      self.RouterPower(router, False, pause_after=0)

  def TurnOffUsedRouters(self):
    """Turns off the routers that were once turned on."""
    for router in self._routers_used:
      self.RouterPower(router, False, pause_after=0)

    self._routers_used = set()


class PyNetworkUITest(pyauto.PyUITest):
  """A subclass of PyUITest for Chrome OS network tests.

  A subclass of PyUITest that automatically sets the flimflam
  priorities to put wifi connections first before starting tests.
  This is for convenience when writing wifi tests.
  """
  _ROUTER_CONFIG_FILE = os.path.join(pyauto.PyUITest.DataDir(),
                                     'pyauto_private', 'chromeos', 'network',
                                     'wifi_testbed_config')

  _FLIMFLAM_PATH = 'org.chromium.flimflam'
  _proxy = dbus.SystemBus().get_object(_FLIMFLAM_PATH, '/')
  _manager = dbus.Interface(_proxy, _FLIMFLAM_PATH + '.Manager')

  def setUp(self):
    # Move ethernet to the end of the flimflam priority list,
    # effectively hiding any ssh connections that the
    # test harness might be using and putting wifi ahead.
    self._PushServiceOrder('vpn,bluetooth,wifi,wimax,cellular,ethernet')
    self._ParseDefaultRoutingTable()
    pyauto.PyUITest.setUp(self)
    self.ForgetAllRememberedNetworks()
    self._wifi_power_strip = None

  def tearDown(self):
    self.ForgetAllRememberedNetworks()
    pyauto.PyUITest.tearDown(self)
    self._PopServiceOrder()
    if self._wifi_power_strip:
      self._wifi_power_strip.TurnOffUsedRouters()
    # Remove the route entry for the power strip.
    if hasattr(self, 'ps_route_entry'):
      os.system('route del -net %(ipaddress)s gateway %(gateway)s netmask '
                '%(netmask)s dev %(iface)s' % self.ps_route_entry)

  def _ParseDefaultRoutingTable(self):
    """Creates and stores a dictionary of the default routing paths."""
    route_table_headers = ['destination', 'gateway', 'genmask', 'flags',
                           'metric', 'ref', 'use', 'iface']
    routes = os.popen('route -n | egrep "^0.0.0.0"').read()
    routes = [interface.split() for interface in routes.split('\n')][:-1]
    self.default_routes = {}
    for iface in routes:
      self.default_routes[iface[-1]] = dict(zip(route_table_headers, iface))

  def ForgetAllRememberedNetworks(self):
    """Forgets all networks that the device has marked as remembered."""
    for service in self.GetNetworkInfo()['remembered_wifi']:
      self.ForgetWifiNetwork(service)

  def _SetServiceOrder(self, service_order):
    self._manager.SetServiceOrder(service_order)
    # Flimflam throws a dbus exception if device is already disabled.  This
    # is not an error.
    try:
      self._manager.DisableTechnology('wifi')
    except dbus.DBusException as e:
      if 'org.chromium.flimflam.Error.AlreadyDisabled' not in str(e):
        raise e
    self._manager.EnableTechnology('wifi')

  def _PushServiceOrder(self, service_order):
    self._old_service_order = self._manager.GetServiceOrder()
    self._SetServiceOrder(service_order)
    service_order = service_order.split(',')

    # Verify services that are present in both the service_order
    # we set and the one retrieved from device are in the correct order.
    set_service_order = self._manager.GetServiceOrder().split(',')
    common_service = set(service_order) & set(set_service_order)

    service_order = [s for s in service_order if s in common_service]
    set_service_order = [s for s in set_service_order if s in common_service]

    assert service_order == set_service_order, \
        'Flimflam service order not set properly. %s != %s' % \
        (service_order, set_service_order)

  def _PopServiceOrder(self):
    self._SetServiceOrder(self._old_service_order)

    # Verify services that are present in both the service_order
    # we set and the one retrieved from device are in the correct order.
    old_service_order = self._old_service_order.split(',')
    set_service_order = self._manager.GetServiceOrder().split(',')
    common_service = set(old_service_order) & set(set_service_order)

    old_service_order = [s for s in old_service_order if s in common_service]
    set_service_order = [s for s in set_service_order if s in common_service]

    assert old_service_order == set_service_order, \
        'Flimflam service order not set properly. %s != %s' % \
        (old_service_order, set_service_order)

  def _SetupRouteForPowerStrip(self, ipaddress, iface='eth'):
    """Create a route table entry for the power strip."""

    # Assume device has only one interface that is prepended with
    # $iface and use that one.
    try:
      iface = [ key for key in self.default_routes.keys() if iface in key ][0]
    except:
      assert 'Unable to find interface of type %s.' % iface

    self.ps_route_entry = {
      'iface' : iface,
      'gateway' : self.default_routes[iface]['gateway'],
      'netmask' : '255.255.255.255',
      'ipaddress' : ipaddress
    }

    os.system('route add -net %(ipaddress)s gateway %(gateway)s netmask '
              '%(netmask)s dev %(iface)s' % self.ps_route_entry)

    # Verify the route was added.
    assert os.system('route -n | egrep "^%(ipaddress)s[[:space:]]+%(gateway)s'
                     '[[:space:]]+%(netmask)s"' % self.ps_route_entry) == 0, \
                     'Failed to create default route for powerstrip.'

  def InitWifiPowerStrip(self):
    """Initializes the router controller using the specified config file."""

    assert os.path.exists(PyNetworkUITest._ROUTER_CONFIG_FILE), \
           'Router configuration file does not exist.'

    config = pyauto.PyUITest.EvalDataFrom(self._ROUTER_CONFIG_FILE)
    strip_ip, routers = config['strip_ip'], config['routers']

    self._SetupRouteForPowerStrip(strip_ip)
    self._wifi_power_strip = WifiPowerStrip(strip_ip, routers)

    self.RouterPower = self._wifi_power_strip.RouterPower
    self.TurnOffAllRouters = self._wifi_power_strip.TurnOffAllRouters
    self.GetRouterConfig = self._wifi_power_strip.GetRouterConfig

  def WaitUntilWifiNetworkAvailable(self, ssid, timeout=60, is_hidden=False):
    """Waits until the given network is available.

    Routers that are just turned on may take up to 1 minute upon turning them
    on to broadcast their SSID.

    Args:
      ssid: SSID of the service we want to connect to.
      timeout: timeout (in seconds)

    Raises:
      Exception if timeout duration has been hit before wifi router is seen.

    Returns:
      True, when the wifi network is seen within the timout period.
      False, otherwise.
    """
    def _GotWifiNetwork():
      # Returns non-empty array if desired SSID is available.
      try:
        return [wifi for wifi in
                self.NetworkScan().get('wifi_networks', {}).values()
                if wifi.get('name') == ssid]
      except pyauto_errors.JSONInterfaceError:
        # Temporary fix until crosbug.com/14174 is fixed.
        # NetworkScan is only used in updating the list of networks so errors
        # thrown by it are not critical to the results of wifi tests that use
        # this method.
        return False

    # The hidden AP's will always be on, thus we will assume it is ready to
    # connect to.
    if is_hidden:
      return bool(_GotWifiNetwork())

    return self.WaitUntil(_GotWifiNetwork, timeout=timeout, retry_sleep=1)

  def GetConnectedWifi(self):
    """Returns the SSID of the currently connected wifi network.

    Returns:
      The SSID of the connected network or None if we're not connected.
    """
    service_list = self.GetNetworkInfo()
    connected_service_path = service_list.get('connected_wifi')
    if 'wifi_networks' in service_list and \
       connected_service_path in service_list['wifi_networks']:
       return service_list['wifi_networks'][connected_service_path]['name']

  def GetServicePath(self, ssid):
    """Returns the service path associated with an SSID.

    Args:
      ssid: String defining the SSID we are searching for.

    Returns:
      The service path or None if SSID does not exist.
    """
    service_list = self.GetNetworkInfo()
    service_list = service_list.get('wifi_networks', [])
    for service_path, service_obj in service_list.iteritems():
      if service_obj['name'] == ssid:
        return service_path
    return None

  def ConnectToWifiRouter(self, router_name):
    """Connects to a router by name.

    Args:
      router_name: The name of the router that is specified in the
                   configuration file.
    """
    router = self._wifi_power_strip.GetRouterConfig(router_name)
    assert router, 'Router with name %s is not defined ' \
                   'in the router configuration.' % router_name
    security = router.get('security', 'SECURITY_NONE')
    passphrase = router.get('passphrase', '')

    # Branch off the connect calls depending on if the wifi network is hidden
    # or not.
    error_string = None
    if router.get('hidden'):
      error_string = self.ConnectToHiddenWifiNetwork(router['ssid'], security,
                                                     passphrase)
    else:
      service_path = self.GetServicePath(router['ssid'])
      assert service_path, 'Service with SSID %s is not present.' % \
             router['ssid']

      logging.debug('Connecting to router %s.' % router_name)
      error_string = self.ConnectToWifiNetwork(service_path, passphrase)
    return error_string