1 # Copyright 2017 - 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 15 import shellescape 16 import signal 17 import time 18 19 from utils import job 20 21 22 class 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 32 def __init__(self, runner, working_dir=None): 33 """Creates a new shell command invoker. 34 35 Args: 36 runner: The object that will run the shell commands. 37 working_dir: The directory that all commands should work in, 38 if none then the runners environment default is used. 39 """ 40 self._runner = runner 41 self._working_dir = working_dir 42 43 def run(self, command, timeout=3600, ignore_status=False): 44 """Runs a generic command through the runner. 45 46 Takes the command and prepares it to be run in the target shell using 47 this objects settings. 48 49 Args: 50 command: The command to run. 51 timeout: How long to wait for the command (in seconds). 52 ignore_status: Whether or not to throw exception based upon status. 53 Returns: 54 A CmdResult object containing the results of the shell command. 55 56 Raises: 57 job.Error: When the command executed but had an error. 58 """ 59 if self._working_dir: 60 command_str = 'cd %s; %s' % (self._working_dir, command) 61 else: 62 command_str = command 63 64 return self._runner.run( 65 command_str, timeout=timeout, ignore_status=ignore_status) 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 environment. 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_command_pids(self, identifier): 95 """Gets the pids of a program, based only upon the starting command 96 97 Searches for a program with a specific name and grabs the pids for all 98 programs that match only the command that started it. The arguments of 99 the command are ignored. 100 101 Args: 102 identifier: The search term that identifies the program. 103 104 Returns: 105 An array of ints of all pids that matched the identifier. 106 """ 107 result = self.run( 108 'ps -C %s --no-heading -o pid:1' % identifier, ignore_status=True) 109 110 # Output looks like pids on separate lines: 111 # 1245 112 # 5001 113 114 pids = result.stdout.splitlines() 115 return [int(pid) for pid in result.stdout.splitlines()] 116 117 def get_pids(self, identifier): 118 """Gets the pids of a program. 119 120 Searches for a program with a specific name and grabs the pids for all 121 programs that match. 122 123 Args: 124 identifier: A search term that identifies the program. 125 126 Returns: An array of all pids that matched the identifier, or None 127 if no pids were found. 128 """ 129 try: 130 result = self.run('ps aux | grep -v grep | grep %s' % identifier) 131 except job.Error: 132 raise StopIteration 133 134 lines = result.stdout.splitlines() 135 136 # The expected output of the above command is like so: 137 # bob 14349 0.0 0.0 34788 5552 pts/2 Ss Oct10 0:03 bash 138 # bob 52967 0.0 0.0 34972 5152 pts/4 Ss Oct10 0:00 bash 139 # Where the format is: 140 # USER PID ... 141 for line in lines: 142 pieces = line.split() 143 yield int(pieces[1]) 144 145 def search_file(self, search_string, file_name): 146 """Searches through a file for a string. 147 148 Args: 149 search_string: The string or pattern to look for. 150 file_name: The name of the file to search. 151 152 Returns: 153 True if the string or pattern was found, False otherwise. 154 """ 155 try: 156 self.run('grep %s %s' % (shellescape.quote(search_string), 157 file_name)) 158 return True 159 except job.Error: 160 return False 161 162 def read_file(self, file_name): 163 """Reads a file through the shell. 164 165 Args: 166 file_name: The name of the file to read. 167 168 Returns: 169 A string of the files contents. 170 """ 171 return self.run('cat %s' % file_name).stdout 172 173 def write_file(self, file_name, data): 174 """Writes a block of data to a file through the shell. 175 176 Args: 177 file_name: The name of the file to write to. 178 data: The string of data to write. 179 """ 180 return self.run('echo %s > %s' % (shellescape.quote(data), file_name)) 181 182 def append_file(self, file_name, data): 183 """Appends a block of data to a file through the shell. 184 185 Args: 186 file_name: The name of the file to write to. 187 data: The string of data to write. 188 """ 189 return self.run('echo %s >> %s' % (shellescape.quote(data), file_name)) 190 191 def touch_file(self, file_name): 192 """Creates a file through the shell. 193 194 Args: 195 file_name: The name of the file to create. 196 """ 197 self.write_file(file_name, '') 198 199 def delete_file(self, file_name): 200 """Deletes a file through the shell. 201 202 Args: 203 file_name: The name of the file to delete. 204 """ 205 try: 206 self.run('rm %s' % file_name) 207 except job.Error as e: 208 if 'No such file or directory' in e.result.stderr: 209 return 210 211 raise 212 213 def kill(self, identifier, timeout=10): 214 """Kills a program or group of programs through the shell. 215 216 Kills all programs that match an identifier through the shell. This 217 will send an increasing queue of kill signals to all programs 218 that match the identifier until either all are dead or the timeout 219 finishes. 220 221 Programs are guaranteed to be killed after running this command. 222 223 Args: 224 identifier: A string used to identify the program. 225 timeout: The time to wait for all programs to die. Each signal will 226 take an equal portion of this time. 227 """ 228 if isinstance(identifier, int): 229 pids = [identifier] 230 else: 231 pids = list(self.get_pids(identifier)) 232 233 signal_queue = [signal.SIGINT, signal.SIGTERM, signal.SIGKILL] 234 235 signal_duration = timeout / len(signal_queue) 236 for sig in signal_queue: 237 for pid in pids: 238 try: 239 self.signal(pid, sig) 240 except job.Error: 241 pass 242 243 start_time = time.time() 244 while pids and time.time() - start_time < signal_duration: 245 time.sleep(0.1) 246 pids = [pid for pid in pids if self.is_alive(pid)] 247 248 if not pids: 249 break 250 251 def signal(self, pid, sig): 252 """Sends a specific signal to a program. 253 254 Args: 255 pid: The process id of the program to kill. 256 sig: The signal to send. 257 258 Raises: 259 job.Error: Raised when the signal fail to reach 260 the specified program. 261 """ 262 self.run('kill -%d %d' % (sig, pid)) 263
