1# Copyright 2016 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import shellescape
16import signal
17import time
18
19from acts.controllers.utils_lib.ssh import connection
20from acts.libs.proc import job
21
22
23class ShellCommand(object):
24    """Wraps basic commands that tend to be tied very closely to a shell.
25
26    This class is a wrapper for running basic shell commands through
27    any object that has a run command. Basic shell functionality for managing
28    the system, programs, and files in wrapped within this class.
29
30    Note: At the moment this only works with the ssh runner.
31    """
32
33    def __init__(self, runner, working_dir=None):
34        """Creates a new shell command invoker.
35
36        Args:
37            runner: The object that will run the shell commands.
38            working_dir: The directory that all commands should work in,
39                         if none then the runners enviroment default is used.
40        """
41        self._runner = runner
42        self._working_dir = working_dir
43
44    def run(self, command, timeout=3600):
45        """Runs a generic command through the runner.
46
47        Takes the command and prepares it to be run in the target shell using
48        this objects settings.
49
50        Args:
51            command: The command to run.
52            timeout: How long to wait for the command (in seconds).
53
54        Returns:
55            A CmdResult object containing the results of the shell command.
56
57        Raises:
58            job.Error: When the command executed but had an error.
59        """
60        if self._working_dir:
61            command_str = 'cd %s; %s' % (self._working_dir, command)
62        else:
63            command_str = command
64
65        return self._runner.run(command_str, timeout=timeout)
66
67    def is_alive(self, identifier):
68        """Checks to see if a program is alive.
69
70        Checks to see if a program is alive on the shells enviroment. This can
71        be used to check on generic programs, or a specific program using
72        a pid.
73
74        Args:
75            identifier: string or int, Used to identify the program to check.
76                        if given an int then it is assumed to be a pid. If
77                        given a string then it will be used as a search key
78                        to compare on the running processes.
79        Returns:
80            True if a process was found running, false otherwise.
81        """
82        try:
83            if isinstance(identifier, str):
84                self.run('ps aux | grep -v grep | grep %s' % identifier)
85            elif isinstance(identifier, int):
86                self.signal(identifier, 0)
87            else:
88                raise ValueError('Bad type was given for identifier')
89
90            return True
91        except job.Error:
92            return False
93
94    def get_pids(self, identifier):
95        """Gets the pids of a program.
96
97        Searches for a program with a specific name and grabs the pids for all
98        programs that match.
99
100        Args:
101            identifier: A search term that identifies the program.
102
103        Returns: An array of all pids that matched the identifier, or None
104                  if no pids were found.
105        """
106        try:
107            result = self.run('ps aux | grep -v grep | grep %s' % identifier)
108        except job.Error:
109            raise StopIteration
110
111        lines = result.stdout.splitlines()
112
113        # The expected output of the above command is like so:
114        # bob    14349  0.0  0.0  34788  5552 pts/2    Ss   Oct10   0:03 bash
115        # bob    52967  0.0  0.0  34972  5152 pts/4    Ss   Oct10   0:00 bash
116        # Where the format is:
117        # USER    PID  ...
118        for line in lines:
119            pieces = line.split()
120            yield int(pieces[1])
121
122    def search_file(self, search_string, file_name):
123        """Searches through a file for a string.
124
125        Args:
126            search_string: The string or pattern to look for.
127            file_name: The name of the file to search.
128
129        Returns:
130            True if the string or pattern was found, False otherwise.
131        """
132        try:
133            self.run('grep %s %s' % (shellescape.quote(search_string),
134                                     file_name))
135            return True
136        except job.Error:
137            return False
138
139    def read_file(self, file_name):
140        """Reads a file through the shell.
141
142        Args:
143            file_name: The name of the file to read.
144
145        Returns:
146            A string of the files contents.
147        """
148        return self.run('cat %s' % file_name).stdout
149
150    def write_file(self, file_name, data):
151        """Writes a block of data to a file through the shell.
152
153        Args:
154            file_name: The name of the file to write to.
155            data: The string of data to write.
156        """
157        return self.run('echo %s > %s' % (shellescape.quote(data), file_name))
158
159    def append_file(self, file_name, data):
160        """Appends a block of data to a file through the shell.
161
162        Args:
163            file_name: The name of the file to write to.
164            data: The string of data to write.
165        """
166        return self.run('echo %s >> %s' % (shellescape.quote(data), file_name))
167
168    def touch_file(self, file_name):
169        """Creates a file through the shell.
170
171        Args:
172            file_name: The name of the file to create.
173        """
174        self.write_file(file_name, '')
175
176    def delete_file(self, file_name):
177        """Deletes a file through the shell.
178
179        Args:
180            file_name: The name of the file to delete.
181        """
182        try:
183            self.run('rm -r %s' % file_name)
184        except job.Error as e:
185            if 'No such file or directory' in e.result.stderr:
186                return
187
188            raise
189
190    def kill(self, identifier, timeout=10):
191        """Kills a program or group of programs through the shell.
192
193        Kills all programs that match an identifier through the shell. This
194        will send an increasing queue of kill signals to all programs
195        that match the identifier until either all are dead or the timeout
196        finishes.
197
198        Programs are guaranteed to be killed after running this command.
199
200        Args:
201            identifier: A string used to identify the program.
202            timeout: The time to wait for all programs to die. Each signal will
203                     take an equal portion of this time.
204        """
205        if isinstance(identifier, int):
206            pids = [identifier]
207        else:
208            pids = list(self.get_pids(identifier))
209
210        signal_queue = [signal.SIGINT, signal.SIGTERM, signal.SIGKILL]
211
212        signal_duration = timeout / len(signal_queue)
213        for sig in signal_queue:
214            for pid in pids:
215                try:
216                    self.signal(pid, sig)
217                except job.Error:
218                    pass
219
220            start_time = time.time()
221            while pids and time.time() - start_time < signal_duration:
222                time.sleep(0.1)
223                pids = [pid for pid in pids if self.is_alive(pid)]
224
225            if not pids:
226                break
227
228    def signal(self, pid, sig):
229        """Sends a specific signal to a program.
230
231        Args:
232            pid: The process id of the program to kill.
233            sig: The signal to send.
234
235        Raises:
236            job.Error: Raised when the signal fail to reach
237                       the specified program.
238        """
239        self.run('kill -%d %d' % (sig, pid))
240