1#!/usr/bin/env python3
2#
3#   Copyright 2022 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0SSHResults
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import subprocess
18import time
19
20from dataclasses import dataclass
21from typing import List, Union
22
23from acts import logger
24from acts import signals
25
26DEFAULT_SSH_USER: str = "fuchsia"
27DEFAULT_SSH_PORT: int = 22
28DEFAULT_SSH_TIMEOUT_SEC: int = 60
29DEFAULT_SSH_CONNECT_TIMEOUT_SEC: int = 30
30DEFAULT_SSH_SERVER_ALIVE_INTERVAL: int = 30
31# The default package repository for all components.
32FUCHSIA_PACKAGE_REPO_NAME = 'fuchsia.com'
33
34
35class SSHResult:
36    """Result of an SSH command."""
37
38    def __init__(
39        self, process: Union[subprocess.CompletedProcess,
40                             subprocess.CalledProcessError]
41    ) -> None:
42        self._raw_stdout = process.stdout
43        self._stdout = process.stdout.decode('utf-8', errors='replace')
44        self._stderr = process.stderr.decode('utf-8', errors='replace')
45        self._exit_status: int = process.returncode
46
47    def __str__(self):
48        if self.exit_status == 0:
49            return self.stdout
50        return f'status {self.exit_status}, stdout: "{self.stdout}", stderr: "{self.stderr}"'
51
52    @property
53    def stdout(self) -> str:
54        return self._stdout
55
56    @property
57    def stderr(self) -> str:
58        return self._stderr
59
60    @property
61    def exit_status(self) -> int:
62        return self._exit_status
63
64    @property
65    def raw_stdout(self) -> bytes:
66        return self._raw_stdout
67
68
69class FuchsiaSSHError(signals.TestError):
70    """A SSH command returned with a non-zero status code."""
71
72    def __init__(self, command: str, result: SSHResult):
73        super().__init__(
74            f'SSH command "{command}" unexpectedly returned {result}')
75        self.result = result
76
77
78class SSHTimeout(signals.TestError):
79    """A SSH command timed out."""
80
81    def __init__(self, err: subprocess.TimeoutExpired):
82        super().__init__(
83            f'SSH command "{err.cmd}" timed out after {err.timeout}s, '
84            f'stdout="{err.stdout}", stderr="{err.stderr}"')
85
86
87class FuchsiaSSHTransportError(signals.TestError):
88    """Failure to send an SSH command."""
89
90
91@dataclass
92class SSHConfig:
93    """SSH client config."""
94
95    # SSH flags. See ssh(1) for full details.
96    host_name: str
97    identity_file: str
98
99    ssh_binary: str = 'ssh'
100    config_file: str = '/dev/null'
101    port: int = 22
102    user: str = DEFAULT_SSH_USER
103
104    # SSH options. See ssh_config(5) for full details.
105    connect_timeout: int = DEFAULT_SSH_CONNECT_TIMEOUT_SEC
106    server_alive_interval: int = DEFAULT_SSH_SERVER_ALIVE_INTERVAL
107    strict_host_key_checking: bool = False
108    user_known_hosts_file: str = "/dev/null"
109    log_level: str = "ERROR"
110
111    def full_command(self, command: str, force_tty: bool = False) -> List[str]:
112        """Generate the complete command to execute command over SSH.
113
114        Args:
115            command: The command to run over SSH
116            force_tty: Force pseudo-terminal allocation. This can be used to
117                execute arbitrary screen-based programs on a remote machine,
118                which can be very useful, e.g. when implementing menu services.
119
120        Returns:
121            Arguments composing the complete call to SSH.
122        """
123        optional_flags = []
124        if force_tty:
125            # Multiple -t options force tty allocation, even if ssh has no local
126            # tty. This is necessary for launching ssh with subprocess without
127            # shell=True.
128            optional_flags.append('-tt')
129
130        return [
131            self.ssh_binary,
132            # SSH flags
133            '-i',
134            self.identity_file,
135            '-F',
136            self.config_file,
137            '-p',
138            str(self.port),
139            # SSH configuration options
140            '-o',
141            f'ConnectTimeout={self.connect_timeout}',
142            '-o',
143            f'ServerAliveInterval={self.server_alive_interval}',
144            '-o',
145            f'StrictHostKeyChecking={"yes" if self.strict_host_key_checking else "no"}',
146            '-o',
147            f'UserKnownHostsFile={self.user_known_hosts_file}',
148            '-o',
149            f'LogLevel={self.log_level}',
150        ] + optional_flags + [
151            f'{self.user}@{self.host_name}'
152        ] + command.split()
153
154
155class SSHProvider:
156    """Device-specific provider for SSH clients."""
157
158    def __init__(self, config: SSHConfig) -> None:
159        """
160        Args:
161            config: SSH client config
162        """
163        logger_tag = f"ssh | {config.host_name}"
164        if config.port != DEFAULT_SSH_PORT:
165            logger_tag += f':{config.port}'
166
167        # Check if the private key exists
168
169        self.log = logger.create_tagged_trace_logger(logger_tag)
170        self.config = config
171
172    def run(self,
173            command: str,
174            timeout_sec: int = DEFAULT_SSH_TIMEOUT_SEC,
175            connect_retries: int = 3,
176            force_tty: bool = False) -> SSHResult:
177        """Run a command on the device then exit.
178
179        Args:
180            command: String to send to the device.
181            timeout_sec: Seconds to wait for the command to complete.
182            connect_retries: Amount of times to retry connect on fail.
183            force_tty: Force pseudo-terminal allocation.
184
185        Raises:
186            FuchsiaSSHError: if the SSH command returns a non-zero status code
187            FuchsiaSSHTimeout: if there is no response within timeout_sec
188            FuchsiaSSHTransportError: if SSH fails to run the command
189
190        Returns:
191            SSHResults from the executed command.
192        """
193        err: Exception
194        for i in range(0, connect_retries):
195            try:
196                return self._run(command, timeout_sec, force_tty)
197            except FuchsiaSSHTransportError as e:
198                err = e
199                self.log.warn(f'Connect failed: {e}')
200        raise err
201
202    def _run(self, command: str, timeout_sec: int, force_tty: bool) -> SSHResult:
203        full_command = self.config.full_command(command, force_tty)
204        self.log.debug(f'Running "{" ".join(full_command)}"')
205        try:
206            process = subprocess.run(full_command,
207                                     capture_output=True,
208                                     timeout=timeout_sec,
209                                     check=True)
210        except subprocess.CalledProcessError as e:
211            if e.returncode == 255:
212                stderr = e.stderr.decode('utf-8', errors='replace')
213                if 'Name or service not known' in stderr or 'Host does not exist' in stderr:
214                    raise FuchsiaSSHTransportError(
215                        f'Hostname {self.config.host_name} cannot be resolved to an address'
216                    ) from e
217                if 'Connection timed out' in stderr:
218                    raise FuchsiaSSHTransportError(
219                        f'Failed to establish a connection to {self.config.host_name} within {timeout_sec}s'
220                    ) from e
221                if 'Connection refused' in stderr:
222                    raise FuchsiaSSHTransportError(
223                        f'Connection refused by {self.config.host_name}') from e
224
225            raise FuchsiaSSHError(command, SSHResult(e)) from e
226        except subprocess.TimeoutExpired as e:
227            raise SSHTimeout(e) from e
228
229        return SSHResult(process)
230
231    def start_v1_component(self,
232                           component: str,
233                           timeout_sec: int = 5,
234                           repo: str = FUCHSIA_PACKAGE_REPO_NAME) -> None:
235        """Start a CFv1 component in the background.
236
237        Args:
238            component: Name of the component without ".cmx".
239            timeout_sec: Seconds to wait for the process to show up in 'ps'.
240            repo: Default package repository for all components.
241
242        Raises:
243            TimeoutError: when the component doesn't launch within timeout_sec
244        """
245        # The "run -d" command will hang when executed without a pseudo-tty
246        # allocated.
247        self.run(
248            f'run -d fuchsia-pkg://{repo}/{component}#meta/{component}.cmx', force_tty=True)
249
250        timeout = time.perf_counter() + timeout_sec
251        while True:
252            ps_cmd = self.run("ps")
253            if f'{component}.cmx' in ps_cmd.stdout:
254                return
255            if time.perf_counter() > timeout:
256                raise TimeoutError(
257                    f'Failed to start "{component}.cmx" after {timeout_sec}s')
258
259    def stop_v1_component(self, component: str) -> None:
260        """Stop all instances of a CFv1 component.
261
262        Args:
263            component: Name of the component without ".cmx"
264        """
265        try:
266            self.run(f'killall {component}.cmx')
267        except FuchsiaSSHError as e:
268            if 'no tasks found' in e.result.stderr:
269                return
270            raise e
271