• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Lint as: python2, python3
2# Copyright 2021 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5"""Autotest communcations to the Hosts (DUTs) via TLS ExecDutCommand."""
6
7import common
8import grpc
9import logging
10import six
11import time
12
13from autotest_lib.server.hosts.tls_client import autotest_common_pb2
14
15from autotest_lib.client.common_lib import error
16from autotest_lib.client.common_lib import utils
17
18
19class TLSExecDutCommandClient():
20    """Object for sending commands to a host, and getting the response."""
21
22    def __init__(self, tlsconnection, hostname):
23        """Configure the grpc channel."""
24        if tlsconnection.alive:
25            self.stub = tlsconnection.stub
26        else:
27            raise error.TLSConnectionError(
28                "TLS connection is not alive when try to creating"
29                " exec_dut_command client.")
30
31        self.hostname = hostname
32        self.tlsconnection = tlsconnection
33
34    def run_cmd(self,
35                cmd,
36                timeout=120,
37                stdout_tee=None,
38                stderr_tee=None,
39                ignore_timeout=False):
40        """
41        Run a command on the host configured during init.
42
43        @param cmd: shell cmd to execute on the DUT
44        @param: stdout_tee/stderr_tee: objects to write the data from the
45            respective streams to
46        @param timeout int(seconds): how long to allow the command to run
47            before forcefully killing it.
48        @param ignore_timeout: if True, do not raise err on timeouts.
49        """
50        if not self.tlsconnection.alive:
51            error.TLSConnectionError(
52                "TLS connection is not up when try to run exec_dut_command.")
53        result = utils.CmdResult(command=cmd)
54        try:
55            self._run(cmd, stdout_tee, stderr_tee, result, timeout)
56        except grpc.RpcError as e:
57            if e.code().name == "DEADLINE_EXCEEDED":
58                if ignore_timeout:
59                    return None
60                raise error.CmdTimeoutError(
61                        cmd, result,
62                        "Command(s) did not complete within %d seconds" %
63                        timeout)
64            raise e
65        except Exception as e:
66            raise e
67        return result
68
69    def _run(self, cmd, stdout_tee, stderr_tee, result, timeout):
70        """Run the provided cmd, populate the result in place."""
71        start_time = time.time()
72        response = self._send_cmd(cmd, timeout)
73
74        stdout_buf = six.StringIO()
75        stderr_buf = six.StringIO()
76        last_status = 0
77
78        if response:
79            for item in response:
80                last_status = item.exit_info.status
81                _log_item(item.stdout, stdout_buf, stdout_tee)
82                _log_item(item.stderr, stderr_buf, stderr_tee)
83
84        result.stdout = stdout_buf.getvalue()
85        result.stderr = stderr_buf.getvalue()
86        result.exit_status = last_status
87        result.duration = time.time() - start_time
88
89    def _send_cmd(self, cmd, timeout):
90        """Serialize and send the cmd to the TLS service."""
91        formatted_cmd = autotest_common_pb2.ExecDutCommandRequest(
92                name=self.hostname, command=cmd)
93        return self.stub.ExecDutCommand(formatted_cmd, timeout=timeout)
94
95
96def _log_item(item, buf, tee):
97    """
98    Parse the provided item.
99
100    If the item exists, append the provided arr with the item & write to
101        the provided tee if provided.
102
103    """
104    if not item:
105        return
106    # TODO dbeckett@ (crbug.com/990593), adjust this to be PY3 compatible.
107    buf.write(item)
108    if tee is not None and tee is not utils.TEE_TO_LOGS:
109        tee.write(item)
110