1# Copyright 2015 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"""This class defines the TestStationHost class."""
6
7import logging
8import os
9
10import common
11
12from autotest_lib.client.bin import local_host
13from autotest_lib.client.common_lib import error
14from autotest_lib.client.common_lib.cros import retry
15from autotest_lib.client.cros import constants as cros_constants
16from autotest_lib.server.hosts import base_classes
17from autotest_lib.server.hosts import moblab_host
18from autotest_lib.server.hosts import ssh_host
19
20
21# TODO(kevcheng): Update the creation method so it's not a research project
22# determining the class inheritance model (same for factory.create_host).
23def create_teststationhost(hostname, **kwargs):
24    """Creates the TestStationHost object.
25
26    @param hostname: Hostname of the test station.
27    @param kwargs: Keyword args to pass to the testbed initialization.
28
29    @return: A Test Station Host object.
30    """
31    classes = [TestStationHost]
32    if hostname == 'localhost':
33        classes.append(local_host.LocalHost)
34    else:
35        classes.append(ssh_host.SSHHost)
36    host_class = type('new_teststationhost', tuple(classes), {})
37    return host_class(hostname, **kwargs)
38
39
40class TestStationHost(base_classes.Host):
41    """This class represents a linux box accessible via ssh."""
42
43
44    def check_credentials(self, hostname):
45        """Make sure teststation credentials work if we're doing ssh.
46
47        @param hostname: Hostname of the machine.
48        """
49        if hostname != 'localhost':
50            try:
51                self.run('true')
52            except error.AutoservRunError:
53                # Some test stations may not have root access, try user adb.
54                logging.debug('Switching to user adb.')
55                self.user = 'adb'
56
57
58    def _initialize(self, hostname='localhost', *args, **dargs):
59        """Initialize a Test Station Host.
60
61        This will create a Test Station Host. Hostname should always refer
62        to the host machine connected to the devices under test.
63
64        @param hostname: Hostname of the machine, default to localhost.
65        """
66        logging.debug('Initializing Test Station Host running on host: %s.',
67                      hostname)
68
69        # Do parent class initializations.
70        super(TestStationHost, self)._initialize(hostname=hostname, *args,
71                                                 **dargs)
72
73        self.check_credentials(hostname)
74
75        # We'll want to do certain things differently if we're on a moblab.
76        self._is_host_moblab = None
77        # Keep track of whether the host was closed since multiple AdbHost
78        # might have an instance of this teststation.
79        self._is_closed = False
80
81
82    @property
83    def is_moblab(self):
84        """Check if the host running adb command is a Moblab.
85
86        @return: True if the host running adb command is a Moblab, False
87                 otherwise.
88        """
89        if self._is_host_moblab is None:
90            try:
91                self.run('cat %s | grep -q moblab' % cros_constants.LSB_RELEASE)
92                self._is_host_moblab = True
93            except (error.AutoservRunError, error.AutotestHostRunError):
94                self._is_host_moblab = False
95        return self._is_host_moblab
96
97
98    def get_tmp_dir(self, parent='/var/tmp'):
99        """Return pathname of a temporary directory on the test station.
100
101        If parent folder is supplied and the teststation is a moblab.  Then
102        the parent will have the moblab tmp directory prepended to it.
103
104        @param parent: The parent dir to create the temporary dir.
105
106        @return: Path of the newly created temporary dir.
107        """
108        if self.is_moblab:
109            parent = (moblab_host.MOBLAB_TMP_DIR if parent == '/tmp'
110                      else os.path.join(moblab_host.MOBLAB_TMP_DIR,
111                                        parent.lstrip('/')))
112        return super(TestStationHost, self).get_tmp_dir(parent=parent)
113
114
115    def run(self, cmd, force_tty=True, *args, **dargs):
116        """Run a command on the adb device.
117
118        This will run the command on the test station.  This method only
119        exists to modify the command supplied if we're running a fastboot
120        command on a moblab, otherwise we leave the command untouched.
121
122        @param cmd: The command line string.
123        @param force_tty: Set to True to force pseudo-terminal allocation to
124                run the command. This allows the command running on remote host
125                to abort when the ssh command is timed out. Default is True.
126
127        @returns A CMDResult object or None if the call timed out and
128                 ignore_timeout is True.
129        """
130        # TODO (sbasi/kevcheng) - Make teststation_host check if running
131        # on Chrome OS, rather than MobLab when prepending sudo to fastboot.
132        if cmd.startswith('fastboot ') and self.is_moblab:
133            cmd = 'sudo -n ' + cmd
134        if force_tty:
135            dargs['options'] = dargs.get('options', '') + ' -t '
136        return super(TestStationHost, self).run(cmd, *args, **dargs)
137
138    @retry.retry(error.GenericHostRunError, timeout_min=10)
139    def download_file(self, src_url, dest_file, unzip=False, unzip_dest=None):
140        """Download the given url.
141
142        @param src_url: The url to download from.
143        @param dest_file: Destination for the file to be downloaded to.
144        @param unzip: If True, unzip the downloaded file.
145        @param unzip_dest: Location to unzip the downloaded file to. If not
146                           provided, dest_file's directory is used.
147
148        @returns: The path of the downloaded file on the teststation.
149        """
150        try:
151            self.run('wget -q -O "%s" "%s"' % (dest_file, src_url))
152
153            readlink_result = self.run('readlink -f "%s"' % dest_file)
154            full_path = readlink_result.stdout.splitlines()[0]
155
156            if unzip:
157                unzip_dest = unzip_dest or os.path.dirname(full_path)
158                self.run('unzip "%s" -x -d "%s"' % (dest_file, unzip_dest))
159
160            return full_path
161        except:
162            # Delete the destination file if download failed.
163            self.run('rm -f "%s"' % dest_file)
164            raise
165