1# Copyright 2016 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"""Utilities used with Brillo hosts."""
6
7import contextlib
8import logging
9
10import common
11from autotest_lib.client.bin import utils
12
13
14_RUN_BACKGROUND_TEMPLATE = '( %(cmd)s ) </dev/null >/dev/null 2>&1 & echo -n $!'
15
16_WAIT_CMD_TEMPLATE = """\
17to=%(timeout)d; \
18while test ${to} -ne 0; do \
19  test $(ps %(pid)d | wc -l) -gt 1 || break; \
20  sleep 1; \
21  to=$((to - 1)); \
22done; \
23test ${to} -ne 0 -o $(ps %(pid)d | wc -l) -eq 1 \
24"""
25
26
27def run_in_background(host, cmd):
28    """Runs a command in the background on the DUT.
29
30    @param host: A host object representing the DUT.
31    @param cmd: The command to run.
32
33    @return The background process ID (integer).
34    """
35    background_cmd = _RUN_BACKGROUND_TEMPLATE % {'cmd': cmd}
36    return int(host.run_output(background_cmd).strip())
37
38
39def wait_for_process(host, pid, timeout=-1):
40    """Waits for a process on the DUT to terminate.
41
42    @param host: A host object representing the DUT.
43    @param pid: The process ID (integer).
44    @param timeout: Number of seconds to wait; default is wait forever.
45
46    @return True if process terminated within the alotted time, False otherwise.
47    """
48    wait_cmd = _WAIT_CMD_TEMPLATE % {'pid': pid, 'timeout': timeout}
49    return host.run(wait_cmd, ignore_status=True).exit_status == 0
50
51
52@contextlib.contextmanager
53def connect_to_ssid(host, ssid, passphrase):
54    """Connects to a given ssid.
55
56    @param host: A host object representing the DUT.
57    @param ssid: A string representing the ssid to connect to. If ssid is None,
58                 assume that host is already connected to wifi.
59    @param passphrase: A string representing the passphrase to the ssid.
60                       Defaults to None.
61    """
62    try:
63        if ssid is None:
64            # No ssid is passed. It is assumed that the host is already
65            # connected to wifi.
66            logging.warning('This test assumes that the device is connected to '
67                            'wifi. If it is not, this test will fail.')
68            yield
69        else:
70            logging.info('Connecting to ssid %s', ssid)
71            # Update the weaved init.rc to stop privet. This sets up shill in
72            # client mode allowing it to connect to wifi.
73            host.remount()
74            host.run('sed \'s/service weaved \/system\/bin\/weaved/'
75                     'service weaved \/system\/bin\/weaved --disable_privet/\' '
76                     '-i /system/etc/init/weaved.rc')
77            host.reboot()
78            utils.poll_for_condition(
79                    lambda: 'running' in host.run('getprop init.svc.shill'
80                                                  ).stdout,
81                    sleep_interval=1, timeout=300,
82                    desc='shill was not started by init')
83            logging.info('Connecting to wifi')
84            wifi_cmd = ('shill_setup_wifi --ssid=%s '
85                        '--wait-for-online-seconds=%i' % (ssid, 300))
86            if passphrase:
87                wifi_cmd += ' --passphrase=%s' % passphrase
88            host.run(wifi_cmd)
89            yield
90    finally:
91        if ssid:
92            # If we connected to a ssid, disconnect.
93            host.remount()
94            host.run('sed \'s/service weaved \/system\/bin\/weaved '
95                     '--disable_privet/service weaved \/system\/bin\/weaved/\' '
96                     '-i /system/etc/init/weaved.rc')
97            host.run('stop weaved')
98            host.run('start weaved')
99