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