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