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