| #!/usr/bin/env python3 |
| # |
| # Copyright 2022 The Fuchsia Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import subprocess |
| import time |
| |
| from dataclasses import dataclass |
| from typing import List, Union |
| |
| from antlion import logger |
| from antlion import signals |
| |
| DEFAULT_SSH_USER: str = "fuchsia" |
| DEFAULT_SSH_PORT: int = 22 |
| DEFAULT_SSH_PRIVATE_KEY: str = "~/.ssh/fuchsia_ed25519" |
| DEFAULT_SSH_TIMEOUT_SEC: int = 60 |
| DEFAULT_SSH_CONNECT_TIMEOUT_SEC: int = 30 |
| DEFAULT_SSH_SERVER_ALIVE_INTERVAL: int = 30 |
| # The default package repository for all components. |
| FUCHSIA_PACKAGE_REPO_NAME = 'fuchsia.com' |
| |
| |
| class SSHResult: |
| """Result of an SSH command.""" |
| |
| def __init__( |
| self, process: Union[subprocess.CompletedProcess, |
| subprocess.CalledProcessError] |
| ) -> None: |
| self._raw_stdout = process.stdout |
| self._stdout = process.stdout.decode('utf-8', errors='replace') |
| self._stderr = process.stderr.decode('utf-8', errors='replace') |
| self._exit_status: int = process.returncode |
| |
| def __str__(self): |
| if self.exit_status == 0: |
| return self.stdout |
| return f'status {self.exit_status}, stdout: "{self.stdout}", stderr: "{self.stderr}"' |
| |
| @property |
| def stdout(self) -> str: |
| return self._stdout |
| |
| @property |
| def stderr(self) -> str: |
| return self._stderr |
| |
| @property |
| def exit_status(self) -> int: |
| return self._exit_status |
| |
| @property |
| def raw_stdout(self) -> bytes: |
| return self._raw_stdout |
| |
| |
| class FuchsiaSSHError(signals.TestError): |
| """A SSH command returned with a non-zero status code.""" |
| |
| def __init__(self, command: str, result: SSHResult): |
| super().__init__( |
| f'SSH command "{command}" unexpectedly returned {result}') |
| self.result = result |
| |
| |
| class SSHTimeout(signals.TestError): |
| """A SSH command timed out.""" |
| |
| def __init__(self, err: subprocess.TimeoutExpired): |
| super().__init__( |
| f'SSH command "{err.cmd}" timed out after {err.timeout}s, ' |
| f'stdout="{err.stdout}", stderr="{err.stderr}"') |
| |
| |
| class FuchsiaSSHTransportError(signals.TestError): |
| """Failure to send an SSH command.""" |
| |
| |
| @dataclass |
| class SSHConfig: |
| """SSH client config.""" |
| |
| # SSH flags. See ssh(1) for full details. |
| host_name: str |
| identity_file: str |
| |
| ssh_binary: str = 'ssh' |
| config_file: str = '/dev/null' |
| port: int = 22 |
| user: str = DEFAULT_SSH_USER |
| |
| # SSH options. See ssh_config(5) for full details. |
| connect_timeout: int = DEFAULT_SSH_CONNECT_TIMEOUT_SEC |
| server_alive_interval: int = DEFAULT_SSH_SERVER_ALIVE_INTERVAL |
| strict_host_key_checking: bool = False |
| user_known_hosts_file: str = "/dev/null" |
| log_level: str = "ERROR" |
| |
| def full_command(self, command: str, force_tty: bool = False) -> List[str]: |
| """Generate the complete command to execute command over SSH. |
| |
| Args: |
| command: The command to run over SSH |
| force_tty: Force pseudo-terminal allocation. This can be used to |
| execute arbitrary screen-based programs on a remote machine, |
| which can be very useful, e.g. when implementing menu services. |
| |
| Returns: |
| Arguments composing the complete call to SSH. |
| """ |
| optional_flags = [] |
| if force_tty: |
| # Multiple -t options force tty allocation, even if ssh has no local |
| # tty. This is necessary for launching ssh with subprocess without |
| # shell=True. |
| optional_flags.append('-tt') |
| |
| return [ |
| self.ssh_binary, |
| # SSH flags |
| '-i', |
| self.identity_file, |
| '-F', |
| self.config_file, |
| '-p', |
| str(self.port), |
| # SSH configuration options |
| '-o', |
| f'ConnectTimeout={self.connect_timeout}', |
| '-o', |
| f'ServerAliveInterval={self.server_alive_interval}', |
| '-o', |
| f'StrictHostKeyChecking={"yes" if self.strict_host_key_checking else "no"}', |
| '-o', |
| f'UserKnownHostsFile={self.user_known_hosts_file}', |
| '-o', |
| f'LogLevel={self.log_level}', |
| ] + optional_flags + [ |
| f'{self.user}@{self.host_name}' |
| ] + command.split() |
| |
| |
| class SSHProvider: |
| """Device-specific provider for SSH clients.""" |
| |
| def __init__(self, config: SSHConfig) -> None: |
| """ |
| Args: |
| config: SSH client config |
| """ |
| logger_tag = f"ssh | {config.host_name}" |
| if config.port != DEFAULT_SSH_PORT: |
| logger_tag += f':{config.port}' |
| |
| # Check if the private key exists |
| |
| self.log = logger.create_tagged_trace_logger(logger_tag) |
| self.config = config |
| |
| def run(self, |
| command: str, |
| timeout_sec: int = DEFAULT_SSH_TIMEOUT_SEC, |
| connect_retries: int = 3, |
| force_tty: bool = False) -> SSHResult: |
| """Run a command on the device then exit. |
| |
| Args: |
| command: String to send to the device. |
| timeout_sec: Seconds to wait for the command to complete. |
| connect_retries: Amount of times to retry connect on fail. |
| force_tty: Force pseudo-terminal allocation. |
| |
| Raises: |
| FuchsiaSSHError: if the SSH command returns a non-zero status code |
| FuchsiaSSHTimeout: if there is no response within timeout_sec |
| FuchsiaSSHTransportError: if SSH fails to run the command |
| |
| Returns: |
| SSHResults from the executed command. |
| """ |
| err: Exception |
| for i in range(0, connect_retries): |
| try: |
| return self._run(command, timeout_sec, force_tty) |
| except FuchsiaSSHTransportError as e: |
| err = e |
| self.log.warn(f'Connect failed: {e}') |
| raise err |
| |
| def _run(self, command: str, timeout_sec: int, force_tty: bool) -> SSHResult: |
| full_command = self.config.full_command(command, force_tty) |
| self.log.debug(f'Running "{" ".join(full_command)}"') |
| try: |
| process = subprocess.run(full_command, |
| capture_output=True, |
| timeout=timeout_sec, |
| check=True) |
| except subprocess.CalledProcessError as e: |
| if e.returncode == 255: |
| stderr = e.stderr.decode('utf-8', errors='replace') |
| if 'Name or service not known' in stderr or 'Host does not exist' in stderr: |
| raise FuchsiaSSHTransportError( |
| f'Hostname {self.config.host_name} cannot be resolved to an address' |
| ) from e |
| if 'Connection timed out' in stderr: |
| raise FuchsiaSSHTransportError( |
| f'Failed to establish a connection to {self.config.host_name} within {timeout_sec}s' |
| ) from e |
| if 'Connection refused' in stderr: |
| raise FuchsiaSSHTransportError( |
| f'Connection refused by {self.config.host_name}') from e |
| |
| raise FuchsiaSSHError(command, SSHResult(e)) from e |
| except subprocess.TimeoutExpired as e: |
| raise SSHTimeout(e) from e |
| |
| return SSHResult(process) |
| |
| def start_v1_component(self, |
| component: str, |
| timeout_sec: int = 5, |
| repo: str = FUCHSIA_PACKAGE_REPO_NAME) -> None: |
| """Start a CFv1 component in the background. |
| |
| Args: |
| component: Name of the component without ".cmx". |
| timeout_sec: Seconds to wait for the process to show up in 'ps'. |
| repo: Default package repository for all components. |
| |
| Raises: |
| TimeoutError: when the component doesn't launch within timeout_sec |
| """ |
| # The "run -d" command will hang when executed without a pseudo-tty |
| # allocated. |
| self.run( |
| f'run -d fuchsia-pkg://{repo}/{component}#meta/{component}.cmx', force_tty=True) |
| |
| timeout = time.perf_counter() + timeout_sec |
| while True: |
| ps_cmd = self.run("ps") |
| if f'{component}.cmx' in ps_cmd.stdout: |
| return |
| if time.perf_counter() > timeout: |
| raise TimeoutError( |
| f'Failed to start "{component}.cmx" after {timeout_sec}s') |
| |
| def stop_v1_component(self, component: str) -> None: |
| """Stop all instances of a CFv1 component. |
| |
| Args: |
| component: Name of the component without ".cmx" |
| """ |
| try: |
| self.run(f'killall {component}.cmx') |
| except FuchsiaSSHError as e: |
| if 'no tasks found' in e.result.stderr: |
| return |
| raise e |
