• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Lint as: python2, python3
2"""Client for Autotest side communcations to the TLS SSH Server."""
3
4
5import grpc
6import logging
7import six
8import time
9
10import common_pb2
11import common_pb2_grpc
12
13from autotest_lib.client.common_lib import error
14from autotest_lib.client.common_lib import utils
15
16TLS_PORT = 7152
17TLS_IP = '10.254.254.254'
18
19class TLSClient(object):
20    """The client side connection to Common-TLS service running in a drone."""
21
22    def __init__(self, hostname):
23        """Configure the grpc channel."""
24        self.hostname = hostname
25        self.channel = grpc.insecure_channel('{}:{}'.format(TLS_IP, TLS_PORT))
26        self.stub = common_pb2_grpc.CommonStub(self.channel)
27
28    def __enter__(self):
29        return self
30
31    def __exit__(self, *exc):
32        self.close()
33
34    def run_cmd(self, cmd, stdout_tee=None, stderr_tee=None, timeout=120):
35        """
36        Run a command on the host configured during init.
37
38        @param cmd: shell cmd to execute on the DUT
39        @param: stdout_tee/stderr_tee: objects to write the data from the
40            respective streams to
41        @param timeout int(seconds): how long to allow the command to run
42            before forcefully killing it.
43
44        """
45        res = utils.CmdResult(command=cmd)
46        try:
47            logging.debug(
48                "Running command %s via TLS ExecDutCommand on host %s",
49                cmd, self.hostname)
50            self._run(cmd, stdout_tee, stderr_tee, res, timeout)
51        except grpc.RpcError as e:
52            if e.code().name == "DEADLINE_EXCEEDED":
53                raise error.CmdTimeoutError(
54                    cmd, res,
55                    "Command(s) did not complete within %d seconds" % timeout)
56            raise e
57        except Exception as e:
58            raise e
59        return res
60
61    def _run(self, cmd, stdout_tee, stderr_tee, res, timeout):
62        """Run the provided cmd, populate the res and return it."""
63        start_time = time.time()
64        response = self._send_cmd(cmd, timeout)
65
66        stdout_buf = six.StringIO()
67        stderr_buf = six.StringIO()
68        last_status = 0
69
70        if response:
71            for item in response:
72                last_status = item.exit_info.status
73                _parse_item_and_log(item.stdout, stdout_buf, stdout_tee)
74                _parse_item_and_log(item.stderr, stderr_buf, stderr_tee)
75
76        res.stdout = stdout_buf.getvalue()
77        res.stderr = stderr_buf.getvalue()
78        res.exit_status = last_status
79        res.duration = time.time() - start_time
80
81    def _send_cmd(self, cmd, timeout):
82        """Serialize and send the cmd to the TLS service."""
83        formatted_cmd = common_pb2.ExecDutCommandRequest(name=self.hostname,
84                                                         command=cmd)
85        return self.stub.ExecDutCommand(formatted_cmd, timeout=timeout)
86
87    def close(self):
88        """Close the grpc channel."""
89        self.channel.close()
90
91
92def _parse_item_and_log(item, buf, tee):
93    """
94    Parse the provided item.
95
96    If the item exists, append the provided arr with the item & write to
97        the provided tee if provided.
98
99    """
100    if not item:
101        return
102    buf.write(item)
103    if tee is not None and tee is not utils.TEE_TO_LOGS:
104        tee.write(item)
105