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        self.re_vt100 = re.compile(r'(\x1b\[|\x9b)[^@-_]*[@-_]|\x1b[@-_]', re.I)
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.encode(errors='replace'))
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 range(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 range(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.output += self.buf[:posafter]
158                    self.buf = self.buf[posafter:]
159                    return earliest_pi
160                tnow_s = time.time()
161                if self.timeout:
162                    tdelta_ms = (tnow_s - tstart_s) * 1000
163                    poll_maxwait = self.timeout - tdelta_ms
164                    if tdelta_ms > self.timeout:
165                        raise Timeout()
166                else:
167                    poll_maxwait = None
168                events = self.poll.poll(poll_maxwait)
169                if not events:
170                    raise Timeout()
171                c = os.read(self.fd, 1024).decode(errors='replace')
172                if not c:
173                    raise EOFError()
174                if self.logfile_read:
175                    self.logfile_read.write(c)
176                self.buf += c
177                # count=0 is supposed to be the default, which indicates
178                # unlimited substitutions, but in practice the version of
179                # Python in Ubuntu 14.04 appears to default to count=2!
180                self.buf = self.re_vt100.sub('', self.buf, count=1000000)
181        finally:
182            if self.logfile_read:
183                self.logfile_read.flush()
184
185    def close(self):
186        """Close the stdio connection to the sub-process.
187
188        This also waits a reasonable time for the sub-process to stop running.
189
190        Args:
191            None.
192
193        Returns:
194            Nothing.
195        """
196
197        os.close(self.fd)
198        for i in range(100):
199            if not self.isalive():
200                break
201            time.sleep(0.1)
202
203    def get_expect_output(self):
204        """Return the output read by expect()
205
206        Returns:
207            The output processed by expect(), as a string.
208        """
209        return self.output
210