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