server/hosts/tls_client/exec_dut_command.py

# Lint as: python2, python3
# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Autotest communcations to the Hosts (DUTs) via TLS ExecDutCommand."""

import common
import grpc
import logging
import six
import time

from autotest_lib.server.hosts.tls_client import autotest_common_pb2

from autotest_lib.client.common_lib import error
from autotest_lib.client.common_lib import utils


class TLSExecDutCommandClient():
    """Object for sending commands to a host, and getting the response."""

    def __init__(self, tlsconnection, hostname):
        """Configure the grpc channel."""
        if tlsconnection.alive:
            self.stub = tlsconnection.stub
        else:
            raise error.TLSConnectionError(
                "TLS connection is not alive when try to creating"
                " exec_dut_command client.")

        self.hostname = hostname
        self.tlsconnection = tlsconnection

    def run_cmd(self,
                cmd,
                timeout=120,
                stdout_tee=None,
                stderr_tee=None,
                ignore_timeout=False):
        """
        Run a command on the host configured during init.

        @param cmd: shell cmd to execute on the DUT
        @param: stdout_tee/stderr_tee: objects to write the data from the
            respective streams to
        @param timeout int(seconds): how long to allow the command to run
            before forcefully killing it.
        @param ignore_timeout: if True, do not raise err on timeouts.
        """
        if not self.tlsconnection.alive:
            error.TLSConnectionError(
                "TLS connection is not up when try to run exec_dut_command.")
        result = utils.CmdResult(command=cmd)
        try:
            self._run(cmd, stdout_tee, stderr_tee, result, timeout)
        except grpc.RpcError as e:
            if e.code().name == "DEADLINE_EXCEEDED":
                if ignore_timeout:
                    return None
                raise error.CmdTimeoutError(
                        cmd, result,
                        "Command(s) did not complete within %d seconds" %
                        timeout)
            raise e
        except Exception as e:
            raise e
        return result

    def _run(self, cmd, stdout_tee, stderr_tee, result, timeout):
        """Run the provided cmd, populate the result in place."""
        start_time = time.time()
        response = self._send_cmd(cmd, timeout)

        stdout_buf = six.StringIO()
        stderr_buf = six.StringIO()
        last_status = 0

        if response:
            for item in response:
                last_status = item.exit_info.status
                _log_item(item.stdout, stdout_buf, stdout_tee)
                _log_item(item.stderr, stderr_buf, stderr_tee)

        result.stdout = stdout_buf.getvalue()
        result.stderr = stderr_buf.getvalue()
        result.exit_status = last_status
        result.duration = time.time() - start_time

    def _send_cmd(self, cmd, timeout):
        """Serialize and send the cmd to the TLS service."""
        formatted_cmd = autotest_common_pb2.ExecDutCommandRequest(
                name=self.hostname, command=cmd)
        return self.stub.ExecDutCommand(formatted_cmd, timeout=timeout)


def _log_item(item, buf, tee):
    """
    Parse the provided item.

    If the item exists, append the provided arr with the item & write to
        the provided tee if provided.

    """
    if not item:
        return
    # TODO dbeckett@ (crbug.com/990593), adjust this to be PY3 compatible.
    buf.write(item)
    if tee is not None and tee is not utils.TEE_TO_LOGS:
        tee.write(item)