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