u_boot_spawn.py 6.3 KB

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