1#
2# Copyright (C) 2016 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16import logging
17import os
18import re
19import shutil
20import tempfile
21
22from vts.runners.host import const
23from vts.utils.python.mirror import mirror_object
24
25
26class ShellMirror(mirror_object.MirrorObject):
27    """The class that acts as the mirror to an Android device's shell terminal.
28
29    Attributes:
30        _client: the TCP client instance.
31        _adb: An AdbProxy object used for interacting with the device via adb.
32        enabled: bool, whether remote shell feature is enabled for the device.
33    """
34
35    TMP_FILE_PATTERN = "/data/local/tmp/nohup.*"
36
37    def __init__(self, client, adb):
38        super(ShellMirror, self).__init__(client)
39        self._adb = adb
40        self.enabled = True
41
42    def Heal(self):
43        """Performs a self healing.
44
45        Includes self diagnosis that looks for any framework errors.
46
47        Returns:
48            bool, True if everything is ok; False otherwise.
49        """
50        res = True
51
52        if self._client:
53            res &= self._client.Heal()
54
55        if not res:
56            logging.error('Self diagnosis found problems in shell mirror.')
57
58        return res
59
60    def Execute(self, command, no_except=False):
61        '''Execute remote shell commands on device.
62
63        Args:
64            command: string or a list of string, shell commands to execute on
65                     device.
66            no_except: bool, if set to True, no exception will be thrown and
67                       error code will be -1 with error message on stderr.
68
69        Returns:
70            A dictionary containing shell command execution results
71        '''
72        if not self.enabled:
73            # TODO(yuexima): use adb shell instead when RPC is disabled
74            return {
75                const.STDOUT: [""] * len(command),
76                const.STDERR:
77                ["VTS remote shell has been disabled."] * len(command),
78                const.EXIT_CODE: [-2] * len(command)
79            }
80        result = self._client.ExecuteShellCommand(command, no_except)
81
82        tmp_dir = tempfile.mkdtemp()
83        pattern = re.compile(self.TMP_FILE_PATTERN)
84
85        for result_val, result_type in zip(
86            [result[const.STDOUT], result[const.STDERR]],
87            ["stdout", "stderr"]):
88            for index, val in enumerate(result_val):
89                # If val is a tmp file name, pull the file and set the contents
90                # to result.
91                if pattern.match(val):
92                    tmp_file = os.path.join(tmp_dir, result_type + str(index))
93                    logging.debug("pulling file: %s to %s", val, tmp_file)
94                    self._adb.pull(val, tmp_file)
95                    result_val[index] = open(tmp_file, "r").read()
96                    self._adb.shell("rm -f %s" % val)
97                else:
98                    result_val[index] = val
99
100        shutil.rmtree(tmp_dir)
101        logging.debug("resp for VTS_AGENT_COMMAND_EXECUTE_SHELL_COMMAND: %s",
102                      result)
103        return result
104
105    def SetConnTimeout(self, timeout):
106        """Set remote shell connection timeout.
107
108        Args:
109            timeout: int, TCP connection timeout in seconds.
110        """
111        self._client.timeout = timeout
112