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