1# Copyright (c) 2013 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
5"""A collection of context managers for working with shill objects."""
6
7import dbus
8import logging
9
10from contextlib import contextmanager
11
12from autotest_lib.client.common_lib import utils
13from autotest_lib.client.cros.networking import shill_proxy
14from autotest_lib.client.cros.power import sys_power
15
16class ContextError(Exception):
17    """An error raised by a context managers dealing with shill objects."""
18    pass
19
20
21class AllowedTechnologiesContext(object):
22    """A context manager for allowing only specified technologies in shill.
23
24    Usage:
25        # Suppose both 'wifi' and 'cellular' technology are originally enabled.
26        allowed = [shill_proxy.ShillProxy.TECHNOLOGY_CELLULAR]
27        with AllowedTechnologiesContext(allowed):
28            # Within this context, only the 'cellular' technology is allowed to
29            # be enabled. The 'wifi' technology is temporarily prohibited and
30            # disabled until after the context ends.
31
32    """
33
34    def __init__(self, allowed):
35        self._allowed = set(allowed)
36
37
38    def __enter__(self):
39        shill = shill_proxy.ShillProxy.get_proxy()
40
41        # The EnabledTechologies property is an array of strings of technology
42        # identifiers.
43        enabled = shill.get_dbus_property(
44                shill.manager,
45                shill_proxy.ShillProxy.MANAGER_PROPERTY_ENABLED_TECHNOLOGIES)
46        self._originally_enabled = set(enabled)
47
48        # The ProhibitedTechnologies property is a comma-separated string of
49        # technology identifiers.
50        prohibited_csv = shill.get_dbus_property(
51                shill.manager,
52                shill_proxy.ShillProxy.MANAGER_PROPERTY_PROHIBITED_TECHNOLOGIES)
53        prohibited = prohibited_csv.split(',') if prohibited_csv else []
54        self._originally_prohibited = set(prohibited)
55
56        prohibited = ((self._originally_prohibited | self._originally_enabled)
57                      - self._allowed)
58        prohibited_csv = ','.join(prohibited)
59
60        logging.debug('Allowed technologies = [%s]', ','.join(self._allowed))
61        logging.debug('Originally enabled technologies = [%s]',
62                      ','.join(self._originally_enabled))
63        logging.debug('Originally prohibited technologies = [%s]',
64                      ','.join(self._originally_prohibited))
65        logging.debug('To be prohibited technologies = [%s]',
66                      ','.join(prohibited))
67
68        # Setting the ProhibitedTechnologies property will disable those
69        # prohibited technologies.
70        shill.set_dbus_property(
71                shill.manager,
72                shill_proxy.ShillProxy.MANAGER_PROPERTY_PROHIBITED_TECHNOLOGIES,
73                prohibited_csv)
74
75        return self
76
77
78    def __exit__(self, exc_type, exc_value, traceback):
79        shill = shill_proxy.ShillProxy.get_proxy()
80
81        prohibited_csv = ','.join(self._originally_prohibited)
82        shill.set_dbus_property(
83                shill.manager,
84                shill_proxy.ShillProxy.MANAGER_PROPERTY_PROHIBITED_TECHNOLOGIES,
85                prohibited_csv)
86
87        # Re-enable originally enabled technologies as they may have been
88        # disabled.
89        enabled = shill.get_dbus_property(
90                shill.manager,
91                shill_proxy.ShillProxy.MANAGER_PROPERTY_ENABLED_TECHNOLOGIES)
92        to_be_reenabled = self._originally_enabled - set(enabled)
93        for technology in to_be_reenabled:
94            shill.manager.EnableTechnology(technology)
95
96        return False
97
98
99class ServiceAutoConnectContext(object):
100    """A context manager for overriding a service's 'AutoConnect' property.
101
102    As the service object of the same service may change during the lifetime
103    of the context, this context manager does not take a service object at
104    construction. Instead, it takes a |get_service| function at construction,
105    which it invokes to obtain a service object when entering and exiting the
106    context. It is assumed that |get_service| always returns a service object
107    that refers to the same service.
108
109    Usage:
110        def get_service():
111            # Some function that returns a service object.
112
113        with ServiceAutoConnectContext(get_service, False):
114            # Within this context, the 'AutoConnect' property of the service
115            # returned by |get_service| is temporarily set to False if it's
116            # initially set to True. The property is restored to its initial
117            # value after the context ends.
118
119    """
120    def __init__(self, get_service, autoconnect):
121        self._get_service = get_service
122        self._autoconnect = autoconnect
123        self._initial_autoconnect = None
124
125
126    def __enter__(self):
127        service = self._get_service()
128        if service is None:
129            raise ContextError('Could not obtain a service object.')
130
131        # Always set the AutoConnect property even if the requested value
132        # is the same so that shill will retain the AutoConnect property, else
133        # shill may override it.
134        service_properties = service.GetProperties()
135        self._initial_autoconnect = shill_proxy.ShillProxy.dbus2primitive(
136            service_properties[
137                shill_proxy.ShillProxy.SERVICE_PROPERTY_AUTOCONNECT])
138        logging.info('ServiceAutoConnectContext: change autoconnect to %s',
139                     self._autoconnect)
140        service.SetProperty(
141            shill_proxy.ShillProxy.SERVICE_PROPERTY_AUTOCONNECT,
142            self._autoconnect)
143
144        # Make sure the cellular service gets persisted by taking it out of
145        # the ephemeral profile.
146        if not service_properties[
147                shill_proxy.ShillProxy.SERVICE_PROPERTY_PROFILE]:
148            shill = shill_proxy.ShillProxy.get_proxy()
149            manager_properties = shill.manager.GetProperties(utf8_strings=True)
150            active_profile = manager_properties[
151                    shill_proxy.ShillProxy.MANAGER_PROPERTY_ACTIVE_PROFILE]
152            logging.info('ServiceAutoConnectContext: change cellular service '
153                         'profile to %s', active_profile)
154            service.SetProperty(
155                    shill_proxy.ShillProxy.SERVICE_PROPERTY_PROFILE,
156                    active_profile)
157
158        return self
159
160
161    def __exit__(self, exc_type, exc_value, traceback):
162        if self._initial_autoconnect != self._autoconnect:
163            service = self._get_service()
164            if service is None:
165                raise ContextError('Could not obtain a service object.')
166
167            logging.info('ServiceAutoConnectContext: restore autoconnect to %s',
168                         self._initial_autoconnect)
169            service.SetProperty(
170                shill_proxy.ShillProxy.SERVICE_PROPERTY_AUTOCONNECT,
171                self._initial_autoconnect)
172        return False
173
174
175    @property
176    def autoconnect(self):
177        """AutoConnect property value within this context."""
178        return self._autoconnect
179
180
181    @property
182    def initial_autoconnect(self):
183        """Initial AutoConnect property value when entering this context."""
184        return self._initial_autoconnect
185
186
187@contextmanager
188def stopped_shill():
189    """A context for executing code which requires shill to be stopped.
190
191    This context stops shill on entry to the context, and starts shill
192    before exit from the context. This context further guarantees that
193    shill will be not restarted by recover_duts and that recover_duts
194    will not otherwise try to run its network connectivity checks, while
195    this context is active.
196
197    Note that the no-restart guarantee applies only if the user of
198    this context completes with a 'reasonable' amount of time. In
199    particular: if networking is unavailable for 15 minutes or more,
200    recover_duts will reboot the DUT.
201
202    """
203    sys_power.pause_check_network_hook()
204    utils.stop_service('shill')
205    yield
206    utils.start_service('shill')
207    sys_power.resume_check_network_hook()
208
209
210class StaticIPContext(object):
211    """StaticIPConfig context manager class.
212
213    Set a StaticIPConfig to the given service.
214
215    """
216    def __init__(self, service, config):
217        self._service = service
218        self._config = config
219
220
221    def __enter__(self):
222        """Configure the StaticIP parameters for the Service and apply those
223        parameters to the interface by forcing a re-connect."""
224        self._service.SetProperty(
225            shill_proxy.ShillProxy.SERVICE_PROPERTY_STATIC_IP_CONFIG,
226            dbus.Dictionary(self._config, signature='sv'))
227        self._service.Disconnect()
228        self._service.Connect()
229
230
231    def __exit__(self, exception, value, traceback):
232        """Clear configuration of StaticIP parameters for the Service and force
233        a re-connect."""
234        self._service.ClearProperty(
235            shill_proxy.ShillProxy.SERVICE_PROPERTY_STATIC_IP_CONFIG)
236        self._service.Disconnect()
237        self._service.Connect()
238