1# Copyright 2016 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import logging
16import os
17import shlex
18import sys
19import time
20
21if os.name == 'posix' and sys.version_info[0] < 3:
22    import subprocess32 as subprocess
23    DEVNULL = open(os.devnull, 'wb')
24else:
25    import subprocess
26    # Only exists in python3.3
27    from subprocess import DEVNULL
28
29
30class Error(Exception):
31    """Indicates that a command failed, is fatal to the test unless caught."""
32
33    def __init__(self, result):
34        super(Error, self).__init__(result)
35        self.result = result
36
37
38class TimeoutError(Error):
39    """Thrown when a BackgroundJob times out on wait."""
40
41
42class Result(object):
43    """Command execution result.
44
45    Contains information on subprocess execution after it has exited.
46
47    Attributes:
48        command: An array containing the command and all arguments that
49                 was executed.
50        exit_status: Integer exit code of the process.
51        stdout_raw: The raw bytes output from standard out.
52        stderr_raw: The raw bytes output from standard error
53        duration: How long the process ran for.
54        did_timeout: True if the program timed out and was killed.
55    """
56
57    @property
58    def stdout(self):
59        """String representation of standard output."""
60        if not self._stdout_str:
61            self._stdout_str = self._raw_stdout.decode(encoding=self._encoding,
62                                                       errors='replace')
63            self._stdout_str = self._stdout_str.strip()
64        return self._stdout_str
65
66    @property
67    def stderr(self):
68        """String representation of standard error."""
69        if not self._stderr_str:
70            self._stderr_str = self._raw_stderr.decode(encoding=self._encoding,
71                                                       errors='replace')
72            self._stderr_str = self._stderr_str.strip()
73        return self._stderr_str
74
75    def __init__(self,
76                 command=[],
77                 stdout=bytes(),
78                 stderr=bytes(),
79                 exit_status=None,
80                 duration=0,
81                 did_timeout=False,
82                 encoding='utf-8'):
83        """
84        Args:
85            command: The command that was run. This will be a list containing
86                     the executed command and all args.
87            stdout: The raw bytes that standard output gave.
88            stderr: The raw bytes that standard error gave.
89            exit_status: The exit status of the command.
90            duration: How long the command ran.
91            did_timeout: True if the command timed out.
92            encoding: The encoding standard that the program uses.
93        """
94        self.command = command
95        self.exit_status = exit_status
96        self._raw_stdout = stdout
97        self._raw_stderr = stderr
98        self._stdout_str = None
99        self._stderr_str = None
100        self._encoding = encoding
101        self.duration = duration
102        self.did_timeout = did_timeout
103
104    def __repr__(self):
105        return ('job.Result(command=%r, stdout=%r, stderr=%r, exit_status=%r, '
106                'duration=%r, did_timeout=%r, encoding=%r)') % (
107                    self.command, self._raw_stdout, self._raw_stderr,
108                    self.exit_status, self.duration, self.did_timeout,
109                    self._encoding)
110
111
112def run(command,
113        timeout=60,
114        ignore_status=False,
115        env=None,
116        io_encoding='utf-8'):
117    """Execute a command in a subproccess and return its output.
118
119    Commands can be either shell commands (given as strings) or the
120    path and arguments to an executable (given as a list).  This function
121    will block until the subprocess finishes or times out.
122
123    Args:
124        command: The command to execute. Can be either a string or a list.
125        timeout: number seconds to wait for command to finish.
126        ignore_status: bool True to ignore the exit code of the remote
127                       subprocess.  Note that if you do ignore status codes,
128                       you should handle non-zero exit codes explicitly.
129        env: dict enviroment variables to setup on the remote host.
130        io_encoding: str unicode encoding of command output.
131
132    Returns:
133        A job.Result containing the results of the ssh command.
134
135    Raises:
136        job.TimeoutError: When the remote command took to long to execute.
137        Error: When the ssh connection failed to be created.
138        CommandError: Ssh worked, but the command had an error executing.
139    """
140    start_time = time.time()
141    proc = subprocess.Popen(
142        command,
143        env=env,
144        stdout=subprocess.PIPE,
145        stderr=subprocess.PIPE,
146        shell=not isinstance(command, list))
147    # Wait on the process terminating
148    timed_out = False
149    out = bytes()
150    err = bytes()
151    try:
152        (out, err) = proc.communicate(timeout=timeout)
153    except subprocess.TimeoutExpired:
154        timed_out = True
155        proc.kill()
156        proc.wait()
157
158    result = Result(
159        command=command,
160        stdout=out,
161        stderr=err,
162        exit_status=proc.returncode,
163        duration=time.time() - start_time,
164        encoding=io_encoding,
165        did_timeout=timed_out)
166    logging.debug(result)
167
168    if timed_out:
169        logging.error("Command %s with %s timeout setting timed out", command,
170                      timeout)
171        raise TimeoutError(result)
172
173    if not ignore_status and proc.returncode != 0:
174        raise Error(result)
175
176    return result
177
178
179def run_async(command, env=None):
180    """Execute a command in a subproccess asynchronously.
181
182    It is the callers responsibility to kill/wait on the resulting
183    subprocess.Popen object.
184
185    Commands can be either shell commands (given as strings) or the
186    path and arguments to an executable (given as a list).  This function
187    will not block.
188
189    Args:
190        command: The command to execute. Can be either a string or a list.
191        env: dict enviroment variables to setup on the remote host.
192
193    Returns:
194        A subprocess.Popen object representing the created subprocess.
195
196    """
197    proc = subprocess.Popen(
198        command,
199        env=env,
200        preexec_fn=os.setpgrp,
201        shell=not isinstance(command, list),
202        stdout=DEVNULL,
203        stderr=subprocess.STDOUT)
204    logging.debug("command %s started with pid %s", command, proc.pid)
205    return proc
206
207