110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
# 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)
|