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"""A module to abstract the shell execution environment on DUT."""
6
7import subprocess
8import tempfile
9
10
11class ShellError(Exception):
12    """Shell specific exception."""
13    pass
14
15
16class LocalShell(object):
17    """An object to wrap the local shell environment."""
18
19    def init(self, os_if):
20        self._os_if = os_if
21
22    def _run_command(self, cmd, block=True):
23        """Helper function of run_command() methods.
24
25        Return the subprocess.Popen() instance to provide access to console
26        output in case command succeeded.  If block=False, will not wait for
27        process to return before returning.
28        """
29        self._os_if.log('Executing %s' % cmd)
30        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,
31                             stderr=subprocess.PIPE)
32        if block:
33            process.wait()
34        return process
35
36    def run_command(self, cmd, block=True):
37        """Run a shell command.
38
39        In case of the command returning an error print its stdout and stderr
40        outputs on the console and dump them into the log. Otherwise suppress
41        all output.
42
43        In case of command error raise an ShellError exception.
44        """
45        process = self._run_command(cmd, block)
46        if process.returncode:
47            err = ['Failed running: %s' % cmd]
48            err.append('stdout:')
49            err.append(process.stdout.read())
50            err.append('stderr:')
51            err.append(process.stderr.read())
52            text = '\n'.join(err)
53            self._os_if.log(text)
54            raise ShellError('command %s failed (code: %d)' %
55                             (cmd, process.returncode))
56
57    def run_command_get_status(self, cmd):
58        """Run a shell command and return its return code.
59
60        The return code of the command is returned, in case of any error.
61        """
62        process = self._run_command(cmd)
63        return process.returncode
64
65    def run_command_get_output(self, cmd):
66        """Run shell command and return its console output to the caller.
67
68        The output is returned as a list of strings stripped of the newline
69        characters.
70        """
71        process = self._run_command(cmd)
72        return [x.rstrip() for x in process.stdout.readlines()]
73
74    def read_file(self, path):
75        """Read the content of the file."""
76        with open(path) as f:
77            return f.read()
78
79    def write_file(self, path, data):
80        """Write the data to the file."""
81        with open(path, 'w') as f:
82            f.write(data)
83
84    def append_file(self, path, data):
85        """Append the data to the file."""
86        with open(path, 'a') as f:
87            f.write(data)
88
89
90class AdbShell(object):
91    """An object to wrap the ADB shell environment.
92
93    DUT is connected to the host in a 1:1 basis. The command is executed
94    via "adb shell".
95    """
96
97    def init(self, os_if):
98        self._os_if = os_if
99        self._host_shell = LocalShell()
100        self._host_shell.init(os_if)
101        self._root_granted = False
102
103    def _run_command(self, cmd):
104        """Helper function of run_command() methods.
105
106        Return the subprocess.Popen() instance to provide access to console
107        output in case command succeeded.
108        """
109        if not self._root_granted:
110            if (self._host_shell.run_command_get_output('adb shell whoami')[0]
111                != 'root'):
112                # Get the root access first as some commands need it.
113                self._host_shell.run_command('adb root')
114            self._root_granted = True
115        cmd = "adb shell '%s'" % cmd.replace("'", "\\'")
116        return self._host_shell._run_command(cmd)
117
118    def run_command(self, cmd):
119        """Run a shell command.
120
121        In case of the command returning an error print its stdout and stderr
122        outputs on the console and dump them into the log. Otherwise suppress
123        all output.
124
125        In case of command error raise an ShellError exception.
126        """
127        process = self._run_command(cmd)
128        if process.returncode:
129            err = ['Failed running: %s' % cmd]
130            err.append('stdout:')
131            err.append(process.stdout.read())
132            err.append('stderr:')
133            err.append(process.stderr.read())
134            text = '\n'.join(err)
135            self._os_if.log(text)
136            raise ShellError('command %s failed (code: %d)' %
137                             (cmd, process.returncode))
138
139    def run_command_get_status(self, cmd):
140        """Run a shell command and return its return code.
141
142        The return code of the command is returned, in case of any error.
143        """
144        # Executing command via adb shell always returns 0.
145        cmd = '(%s); echo $?' % cmd
146        lines = self.run_command_get_output(cmd)
147        if len(lines) == 0:
148            raise ShellError('Somthing wrong on getting status: %r' % lines)
149        return int(lines[-1])
150
151    def run_command_get_output(self, cmd):
152        """Run shell command and return its console output to the caller.
153
154        The output is returned as a list of strings stripped of the newline
155        characters.
156        """
157        # stderr is merged into stdout through adb shell.
158        cmd = '(%s) 2>/dev/null' % cmd
159        process = self._run_command(cmd)
160        return [x.rstrip() for x in process.stdout.readlines()]
161
162    def read_file(self, path):
163        """Read the content of the file."""
164        with tempfile.NamedTemporaryFile() as f:
165            cmd = 'adb pull %s %s' % (path, f.name)
166            self._host_shell.run_command(cmd)
167            return self._host_shell.read_file(f.name)
168
169    def write_file(self, path, data):
170        """Write the data to the file."""
171        with tempfile.NamedTemporaryFile() as f:
172            self._host_shell.write_file(f.name, data)
173            cmd = 'adb push %s %s' % (f.name, path)
174            self._host_shell.run_command(cmd)
175
176    def append_file(self, path, data):
177        """Append the data to the file."""
178        with tempfile.NamedTemporaryFile() as f:
179            cmd = 'adb pull %s %s' % (path, f.name)
180            self._host_shell.run_command(cmd)
181            self._host_shell.append_file(f.name, data)
182            cmd = 'adb push %s %s' % (f.name, path)
183            self._host_shell.run_command(cmd)
184
185    def wait_for_device(self, timeout):
186        """Wait for an Android device connected."""
187        cmd = 'timeout %s adb wait-for-device' % timeout
188        return self._host_shell.run_command_get_status(cmd) == 0
189
190    def wait_for_no_device(self, timeout):
191        """Wait for no Android connected (offline)."""
192        cmd = ('for i in $(seq 0 %d); do adb shell sleep 1 || false; done' %
193               timeout)
194        return self._host_shell.run_command_get_status(cmd) != 0
195