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.cros import constants as cros_constants
15from autotest_lib.server.hosts import base_classes
16from autotest_lib.server.hosts import moblab_host
17from autotest_lib.server.hosts import ssh_host
18
19
20# TODO(kevcheng): Update the creation method so it's not a research project
21# determining the class inheritance model (same for factory.create_host).
22def create_teststationhost(hostname, **kwargs):
23    """Creates the TestStationHost object.
24
25    @param hostname: Hostname of the test station.
26    @param kwargs: Keyword args to pass to the testbed initialization.
27
28    @return: A Test Station Host object.
29    """
30    classes = [TestStationHost]
31    if hostname == 'localhost':
32        classes.append(local_host.LocalHost)
33    else:
34        classes.append(ssh_host.SSHHost)
35    host_class = type('new_teststationhost', tuple(classes), {})
36    return host_class(hostname, **kwargs)
37
38
39class TestStationHost(base_classes.Host):
40    """This class represents a linux box accessible via ssh."""
41
42
43    def check_credentials(self, hostname):
44        """Make sure teststation credentials work if we're doing ssh.
45
46        @param hostname: Hostname of the machine.
47        """
48        if hostname != 'localhost':
49            try:
50                self.run('true')
51            except error.AutoservRunError:
52                # Some test stations may not have root access, try user adb.
53                logging.debug('Switching to user adb.')
54                self.user = 'adb'
55
56
57    def _initialize(self, hostname='localhost', *args, **dargs):
58        """Initialize a Test Station Host.
59
60        This will create a Test Station Host. Hostname should always refer
61        to the host machine connected to the devices under test.
62
63        @param hostname: Hostname of the machine, default to localhost.
64        """
65        logging.debug('Initializing Test Station Host running on host: %s.',
66                      hostname)
67
68        # Do parent class initializations.
69        super(TestStationHost, self)._initialize(hostname=hostname, *args,
70                                                 **dargs)
71
72        self.check_credentials(hostname)
73
74        # We'll want to do certain things differently if we're on a moblab.
75        self._is_host_moblab = None
76        # Keep track of whether the host was closed since multiple AdbHost
77        # might have an instance of this teststation.
78        self._is_closed = False
79
80
81    @property
82    def is_moblab(self):
83        """Check if the host running adb command is a Moblab.
84
85        @return: True if the host running adb command is a Moblab, False
86                 otherwise.
87        """
88        if self._is_host_moblab is None:
89            try:
90                self.run('cat %s | grep -q moblab' % cros_constants.LSB_RELEASE)
91                self._is_host_moblab = True
92            except (error.AutoservRunError, error.AutotestHostRunError):
93                self._is_host_moblab = False
94        return self._is_host_moblab
95
96
97    def get_tmp_dir(self, parent='/tmp'):
98        """Return pathname of a temporary directory on the test station.
99
100        If parent folder is supplied and the teststation is a moblab.  Then
101        the parent will have the moblab tmp directory prepended to it.
102
103        @param parent: The parent dir to create the temporary dir.
104
105        @return: Path of the newly created temporary dir.
106        """
107        if self.is_moblab:
108            parent = (moblab_host.MOBLAB_TMP_DIR if parent == '/tmp'
109                      else os.path.join(moblab_host.MOBLAB_TMP_DIR,
110                                        parent.lstrip('/')))
111        return super(TestStationHost, self).get_tmp_dir(parent=parent)
112
113
114    def run(self, cmd, *args, **dargs):
115        """Run a command on the adb device.
116
117        This will run the command on the test station.  This method only
118        exists to modify the command supplied if we're running a fastboot
119        command on a moblab, otherwise we leave the command untouched.
120
121        @param cmd: The command line string.
122
123        @returns A CMDResult object or None if the call timed out and
124                 ignore_timeout is True.
125        """
126        # TODO (sbasi/kevcheng) - Make teststation_host check if running
127        # on Chrome OS, rather than MobLab when prepending sudo to fastboot.
128        if cmd.startswith('fastboot ') and self.is_moblab:
129            cmd = 'sudo -n ' + cmd
130        return super(TestStationHost, self).run(cmd, *args, **dargs)
131
132
133    def close(self):
134        if not self._is_closed:
135            self._is_closed = True
136            super(TestStationHost, self).close()
137        else:
138            logging.debug('Teststaion already closed.')
139