xref: /openbmc/u-boot/test/py/u_boot_spawn.py (revision c7c47ca2) 
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