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): 24 '''Spawn (fork/exec) the sub-process. 25 26 Args: 27 args: array of processs arguments. argv[0] is the command to execute. 28 29 Returns: 30 Nothing. 31 ''' 32 33 self.waited = False 34 self.buf = '' 35 self.logfile_read = None 36 self.before = '' 37 self.after = '' 38 self.timeout = None 39 40 (self.pid, self.fd) = pty.fork() 41 if self.pid == 0: 42 try: 43 # For some reason, SIGHUP is set to SIG_IGN at this point when 44 # run under "go" (www.go.cd). Perhaps this happens under any 45 # background (non-interactive) system? 46 signal.signal(signal.SIGHUP, signal.SIG_DFL) 47 os.execvp(args[0], args) 48 except: 49 print 'CHILD EXECEPTION:' 50 import traceback 51 traceback.print_exc() 52 finally: 53 os._exit(255) 54 55 self.poll = select.poll() 56 self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL) 57 58 def kill(self, sig): 59 '''Send unix signal "sig" to the child process. 60 61 Args: 62 sig: The signal number to send. 63 64 Returns: 65 Nothing. 66 ''' 67 68 os.kill(self.pid, sig) 69 70 def isalive(self): 71 '''Determine whether the child process is still running. 72 73 Args: 74 None. 75 76 Returns: 77 Boolean indicating whether process is alive. 78 ''' 79 80 if self.waited: 81 return False 82 83 w = os.waitpid(self.pid, os.WNOHANG) 84 if w[0] == 0: 85 return True 86 87 self.waited = True 88 return False 89 90 def send(self, data): 91 '''Send data to the sub-process's stdin. 92 93 Args: 94 data: The data to send to the process. 95 96 Returns: 97 Nothing. 98 ''' 99 100 os.write(self.fd, data) 101 102 def expect(self, patterns): 103 '''Wait for the sub-process to emit specific data. 104 105 This function waits for the process to emit one pattern from the 106 supplied list of patterns, or for a timeout to occur. 107 108 Args: 109 patterns: A list of strings or regex objects that we expect to 110 see in the sub-process' stdout. 111 112 Returns: 113 The index within the patterns array of the pattern the process 114 emitted. 115 116 Notable exceptions: 117 Timeout, if the process did not emit any of the patterns within 118 the expected time. 119 ''' 120 121 for pi in xrange(len(patterns)): 122 if type(patterns[pi]) == type(''): 123 patterns[pi] = re.compile(patterns[pi]) 124 125 try: 126 while True: 127 earliest_m = None 128 earliest_pi = None 129 for pi in xrange(len(patterns)): 130 pattern = patterns[pi] 131 m = pattern.search(self.buf) 132 if not m: 133 continue 134 if earliest_m and m.start() > earliest_m.start(): 135 continue 136 earliest_m = m 137 earliest_pi = pi 138 if earliest_m: 139 pos = earliest_m.start() 140 posafter = earliest_m.end() + 1 141 self.before = self.buf[:pos] 142 self.after = self.buf[pos:posafter] 143 self.buf = self.buf[posafter:] 144 return earliest_pi 145 events = self.poll.poll(self.timeout) 146 if not events: 147 raise Timeout() 148 c = os.read(self.fd, 1024) 149 if not c: 150 raise EOFError() 151 if self.logfile_read: 152 self.logfile_read.write(c) 153 self.buf += c 154 finally: 155 if self.logfile_read: 156 self.logfile_read.flush() 157 158 def close(self): 159 '''Close the stdio connection to the sub-process. 160 161 This also waits a reasonable time for the sub-process to stop running. 162 163 Args: 164 None. 165 166 Returns: 167 Nothing. 168 ''' 169 170 os.close(self.fd) 171 for i in xrange(100): 172 if not self.isalive(): 173 break 174 time.sleep(0.1) 175