1# SPDX-License-Identifier: GPL-2.0 2# Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved. 3 4""" 5Logic to spawn a sub-process and interact with its stdio. 6""" 7 8import os 9import re 10import pty 11import signal 12import select 13import time 14import traceback 15 16class Timeout(Exception): 17 """An exception sub-class that indicates that a timeout occurred.""" 18 19class Spawn: 20 """Represents the stdio of a freshly created sub-process. Commands may be 21 sent to the process, and responses waited for. 22 23 Members: 24 output: accumulated output from expect() 25 """ 26 27 def __init__(self, args, cwd=None): 28 """Spawn (fork/exec) the sub-process. 29 30 Args: 31 args: array of processs arguments. argv[0] is the command to 32 execute. 33 cwd: the directory to run the process in, or None for no change. 34 35 Returns: 36 Nothing. 37 """ 38 39 self.waited = False 40 self.exit_code = 0 41 self.exit_info = '' 42 self.buf = '' 43 self.output = '' 44 self.logfile_read = None 45 self.before = '' 46 self.after = '' 47 self.timeout = None 48 # http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences 49 self.re_vt100 = re.compile(r'(\x1b\[|\x9b)[^@-_]*[@-_]|\x1b[@-_]', re.I) 50 51 (self.pid, self.fd) = pty.fork() 52 if self.pid == 0: 53 try: 54 # For some reason, SIGHUP is set to SIG_IGN at this point when 55 # run under "go" (www.go.cd). Perhaps this happens under any 56 # background (non-interactive) system? 57 signal.signal(signal.SIGHUP, signal.SIG_DFL) 58 if cwd: 59 os.chdir(cwd) 60 os.execvp(args[0], args) 61 except: 62 print('CHILD EXECEPTION:') 63 traceback.print_exc() 64 finally: 65 os._exit(255) 66 67 try: 68 self.poll = select.poll() 69 self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | 70 select.POLLHUP | select.POLLNVAL) 71 except: 72 self.close() 73 raise 74 75 def kill(self, sig): 76 """Send unix signal "sig" to the child process. 77 78 Args: 79 sig: The signal number to send. 80 81 Returns: 82 Nothing. 83 """ 84 85 os.kill(self.pid, sig) 86 87 def checkalive(self): 88 """Determine whether the child process is still running. 89 90 Returns: 91 tuple: 92 True if process is alive, else False 93 0 if process is alive, else exit code of process 94 string describing what happened ('' or 'status/signal n') 95 """ 96 97 if self.waited: 98 return False, self.exit_code, self.exit_info 99 100 w = os.waitpid(self.pid, os.WNOHANG) 101 if w[0] == 0: 102 return True, 0, 'running' 103 status = w[1] 104 105 if os.WIFEXITED(status): 106 self.exit_code = os.WEXITSTATUS(status) 107 self.exit_info = 'status %d' % self.exit_code 108 elif os.WIFSIGNALED(status): 109 signum = os.WTERMSIG(status) 110 self.exit_code = -signum 111 self.exit_info = 'signal %d (%s)' % (signum, signal.Signals(signum).name) 112 self.waited = True 113 return False, self.exit_code, self.exit_info 114 115 def isalive(self): 116 """Determine whether the child process is still running. 117 118 Args: 119 None. 120 121 Returns: 122 Boolean indicating whether process is alive. 123 """ 124 return self.checkalive()[0] 125 126 def send(self, data): 127 """Send data to the sub-process's stdin. 128 129 Args: 130 data: The data to send to the process. 131 132 Returns: 133 Nothing. 134 """ 135 136 os.write(self.fd, data.encode(errors='replace')) 137 138 def expect(self, patterns): 139 """Wait for the sub-process to emit specific data. 140 141 This function waits for the process to emit one pattern from the 142 supplied list of patterns, or for a timeout to occur. 143 144 Args: 145 patterns: A list of strings or regex objects that we expect to 146 see in the sub-process' stdout. 147 148 Returns: 149 The index within the patterns array of the pattern the process 150 emitted. 151 152 Notable exceptions: 153 Timeout, if the process did not emit any of the patterns within 154 the expected time. 155 """ 156 157 for pi in range(len(patterns)): 158 if type(patterns[pi]) == type(''): 159 patterns[pi] = re.compile(patterns[pi]) 160 161 tstart_s = time.time() 162 try: 163 while True: 164 earliest_m = None 165 earliest_pi = None 166 for pi in range(len(patterns)): 167 pattern = patterns[pi] 168 m = pattern.search(self.buf) 169 if not m: 170 continue 171 if earliest_m and m.start() >= earliest_m.start(): 172 continue 173 earliest_m = m 174 earliest_pi = pi 175 if earliest_m: 176 pos = earliest_m.start() 177 posafter = earliest_m.end() 178 self.before = self.buf[:pos] 179 self.after = self.buf[pos:posafter] 180 self.output += self.buf[:posafter] 181 self.buf = self.buf[posafter:] 182 return earliest_pi 183 tnow_s = time.time() 184 if self.timeout: 185 tdelta_ms = (tnow_s - tstart_s) * 1000 186 poll_maxwait = self.timeout - tdelta_ms 187 if tdelta_ms > self.timeout: 188 raise Timeout() 189 else: 190 poll_maxwait = None 191 events = self.poll.poll(poll_maxwait) 192 if not events: 193 raise Timeout() 194 try: 195 c = os.read(self.fd, 1024).decode(errors='replace') 196 except OSError as err: 197 # With sandbox, try to detect when U-Boot exits when it 198 # shouldn't and explain why. This is much more friendly than 199 # just dying with an I/O error 200 if err.errno == 5: # Input/output error 201 alive, _, info = self.checkalive() 202 if alive: 203 raise err 204 raise ValueError('U-Boot exited with %s' % info) 205 raise err 206 if self.logfile_read: 207 self.logfile_read.write(c) 208 self.buf += c 209 # count=0 is supposed to be the default, which indicates 210 # unlimited substitutions, but in practice the version of 211 # Python in Ubuntu 14.04 appears to default to count=2! 212 self.buf = self.re_vt100.sub('', self.buf, count=1000000) 213 finally: 214 if self.logfile_read: 215 self.logfile_read.flush() 216 217 def close(self): 218 """Close the stdio connection to the sub-process. 219 220 This also waits a reasonable time for the sub-process to stop running. 221 222 Args: 223 None. 224 225 Returns: 226 Nothing. 227 """ 228 229 os.close(self.fd) 230 for _ in range(100): 231 if not self.isalive(): 232 break 233 time.sleep(0.1) 234 235 def get_expect_output(self): 236 """Return the output read by expect() 237 238 Returns: 239 The output processed by expect(), as a string. 240 """ 241 return self.output 242