1# SPDX-License-Identifier: GPL-2.0
2# Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
3
4"""
5Logic to spawn a sub-process and interact with its stdio.
6"""
7
8import os
9import re
10import pty
11import signal
12import select
13import time
14import traceback
15
16class Timeout(Exception):
17    """An exception sub-class that indicates that a timeout occurred."""
18
19class Spawn:
20    """Represents the stdio of a freshly created sub-process. Commands may be
21    sent to the process, and responses waited for.
22
23    Members:
24        output: accumulated output from expect()
25    """
26
27    def __init__(self, args, cwd=None):
28        """Spawn (fork/exec) the sub-process.
29
30        Args:
31            args: array of processs arguments. argv[0] is the command to
32              execute.
33            cwd: the directory to run the process in, or None for no change.
34
35        Returns:
36            Nothing.
37        """
38
39        self.waited = False
40        self.exit_code = 0
41        self.exit_info = ''
42        self.buf = ''
43        self.output = ''
44        self.logfile_read = None
45        self.before = ''
46        self.after = ''
47        self.timeout = None
48        # http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences
49        self.re_vt100 = re.compile(r'(\x1b\[|\x9b)[^@-_]*[@-_]|\x1b[@-_]', re.I)
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                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 |
70                               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 checkalive(self):
88        """Determine whether the child process is still running.
89
90        Returns:
91            tuple:
92                True if process is alive, else False
93                0 if process is alive, else exit code of process
94                string describing what happened ('' or 'status/signal n')
95        """
96
97        if self.waited:
98            return False, self.exit_code, self.exit_info
99
100        w = os.waitpid(self.pid, os.WNOHANG)
101        if w[0] == 0:
102            return True, 0, 'running'
103        status = w[1]
104
105        if os.WIFEXITED(status):
106            self.exit_code = os.WEXITSTATUS(status)
107            self.exit_info = 'status %d' % self.exit_code
108        elif os.WIFSIGNALED(status):
109            signum = os.WTERMSIG(status)
110            self.exit_code = -signum
111            self.exit_info = 'signal %d (%s)' % (signum, signal.Signals(signum).name)
112        self.waited = True
113        return False, self.exit_code, self.exit_info
114
115    def isalive(self):
116        """Determine whether the child process is still running.
117
118        Args:
119            None.
120
121        Returns:
122            Boolean indicating whether process is alive.
123        """
124        return self.checkalive()[0]
125
126    def send(self, data):
127        """Send data to the sub-process's stdin.
128
129        Args:
130            data: The data to send to the process.
131
132        Returns:
133            Nothing.
134        """
135
136        os.write(self.fd, data.encode(errors='replace'))
137
138    def expect(self, patterns):
139        """Wait for the sub-process to emit specific data.
140
141        This function waits for the process to emit one pattern from the
142        supplied list of patterns, or for a timeout to occur.
143
144        Args:
145            patterns: A list of strings or regex objects that we expect to
146                see in the sub-process' stdout.
147
148        Returns:
149            The index within the patterns array of the pattern the process
150            emitted.
151
152        Notable exceptions:
153            Timeout, if the process did not emit any of the patterns within
154            the expected time.
155        """
156
157        for pi in range(len(patterns)):
158            if type(patterns[pi]) == type(''):
159                patterns[pi] = re.compile(patterns[pi])
160
161        tstart_s = time.time()
162        try:
163            while True:
164                earliest_m = None
165                earliest_pi = None
166                for pi in range(len(patterns)):
167                    pattern = patterns[pi]
168                    m = pattern.search(self.buf)
169                    if not m:
170                        continue
171                    if earliest_m and m.start() >= earliest_m.start():
172                        continue
173                    earliest_m = m
174                    earliest_pi = pi
175                if earliest_m:
176                    pos = earliest_m.start()
177                    posafter = earliest_m.end()
178                    self.before = self.buf[:pos]
179                    self.after = self.buf[pos:posafter]
180                    self.output += self.buf[:posafter]
181                    self.buf = self.buf[posafter:]
182                    return earliest_pi
183                tnow_s = time.time()
184                if self.timeout:
185                    tdelta_ms = (tnow_s - tstart_s) * 1000
186                    poll_maxwait = self.timeout - tdelta_ms
187                    if tdelta_ms > self.timeout:
188                        raise Timeout()
189                else:
190                    poll_maxwait = None
191                events = self.poll.poll(poll_maxwait)
192                if not events:
193                    raise Timeout()
194                try:
195                    c = os.read(self.fd, 1024).decode(errors='replace')
196                except OSError as err:
197                    # With sandbox, try to detect when U-Boot exits when it
198                    # shouldn't and explain why. This is much more friendly than
199                    # just dying with an I/O error
200                    if err.errno == 5:  # Input/output error
201                        alive, _, info = self.checkalive()
202                        if alive:
203                            raise err
204                        raise ValueError('U-Boot exited with %s' % info)
205                    raise err
206                if self.logfile_read:
207                    self.logfile_read.write(c)
208                self.buf += c
209                # count=0 is supposed to be the default, which indicates
210                # unlimited substitutions, but in practice the version of
211                # Python in Ubuntu 14.04 appears to default to count=2!
212                self.buf = self.re_vt100.sub('', self.buf, count=1000000)
213        finally:
214            if self.logfile_read:
215                self.logfile_read.flush()
216
217    def close(self):
218        """Close the stdio connection to the sub-process.
219
220        This also waits a reasonable time for the sub-process to stop running.
221
222        Args:
223            None.
224
225        Returns:
226            Nothing.
227        """
228
229        os.close(self.fd)
230        for _ in range(100):
231            if not self.isalive():
232                break
233            time.sleep(0.1)
234
235    def get_expect_output(self):
236        """Return the output read by expect()
237
238        Returns:
239            The output processed by expect(), as a string.
240        """
241        return self.output
242