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