1#!/usr/bin/env python3
2#
3#   Copyright 2019 - 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.0
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 backoff
18import os
19import logging
20import paramiko
21import socket
22import time
23
24from acts import utils
25from acts.controllers.fuchsia_lib.base_lib import DeviceOffline
26
27logging.getLogger("paramiko").setLevel(logging.WARNING)
28# paramiko-ng will throw INFO messages when things get disconnect or cannot
29# connect perfectly the first time.  In this library those are all handled by
30# either retrying and/or throwing an exception for the appropriate case.
31# Therefore, in order to reduce confusion in the logs the log level is set to
32# WARNING.
33
34
35def get_private_key(ip_address, ssh_config):
36    """Tries to load various ssh key types.
37
38    Args:
39        ip_address: IP address of ssh server.
40        ssh_config: ssh_config location for the ssh server.
41    Returns:
42        The ssh private key
43    """
44    exceptions = []
45    try:
46        logging.debug('Trying to load SSH key type: ed25519')
47        return paramiko.ed25519key.Ed25519Key(
48            filename=get_ssh_key_for_host(ip_address, ssh_config))
49    except paramiko.SSHException as e:
50        exceptions.append(e)
51        logging.debug('Failed loading SSH key type: ed25519')
52
53    try:
54        logging.debug('Trying to load SSH key type: rsa')
55        return paramiko.RSAKey.from_private_key_file(
56            filename=get_ssh_key_for_host(ip_address, ssh_config))
57    except paramiko.SSHException as e:
58        exceptions.append(e)
59        logging.debug('Failed loading SSH key type: rsa')
60
61    raise Exception('No valid ssh key type found', exceptions)
62
63
64@backoff.on_exception(
65    backoff.constant,
66    (paramiko.ssh_exception.SSHException,
67     paramiko.ssh_exception.AuthenticationException, socket.timeout,
68     socket.error, ConnectionRefusedError, ConnectionResetError),
69    interval=1.5,
70    max_tries=4)
71def create_ssh_connection(ip_address,
72                          ssh_username,
73                          ssh_config,
74                          connect_timeout=10,
75                          auth_timeout=10,
76                          banner_timeout=10):
77    """Creates and ssh connection to a Fuchsia device
78
79    Args:
80        ip_address: IP address of ssh server.
81        ssh_username: Username for ssh server.
82        ssh_config: ssh_config location for the ssh server.
83        connect_timeout: Timeout value for connecting to ssh_server.
84        auth_timeout: Timeout value to wait for authentication.
85        banner_timeout: Timeout to wait for ssh banner.
86
87    Returns:
88        A paramiko ssh object
89    """
90    if not utils.is_pingable(ip_address):
91        raise DeviceOffline("Device %s is not reachable via "
92                            "the network." % ip_address)
93    ssh_key = get_private_key(ip_address=ip_address, ssh_config=ssh_config)
94    ssh_client = paramiko.SSHClient()
95    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
96    ssh_client.connect(hostname=ip_address,
97                       username=ssh_username,
98                       allow_agent=False,
99                       pkey=ssh_key,
100                       timeout=connect_timeout,
101                       auth_timeout=auth_timeout,
102                       banner_timeout=banner_timeout)
103    ssh_client.get_transport().set_keepalive(1)
104    return ssh_client
105
106
107def ssh_is_connected(ssh_client):
108    """Checks to see if the SSH connection is alive.
109    Args:
110        ssh_client: A paramiko SSH client instance.
111    Returns:
112          True if connected, False or None if not connected.
113    """
114    return ssh_client and ssh_client.get_transport().is_active()
115
116
117def get_ssh_key_for_host(host, ssh_config_file):
118    """Gets the SSH private key path from a supplied ssh_config_file and the
119       host.
120    Args:
121        host (str): The ip address or host name that SSH will connect to.
122        ssh_config_file (str): Path to the ssh_config_file that will be used
123            to connect to the host.
124
125    Returns:
126        path: A path to the private key for the SSH connection.
127    """
128    ssh_config = paramiko.SSHConfig()
129    user_config_file = os.path.expanduser(ssh_config_file)
130    if os.path.exists(user_config_file):
131        with open(user_config_file) as f:
132            ssh_config.parse(f)
133    user_config = ssh_config.lookup(host)
134
135    if 'identityfile' not in user_config:
136        raise ValueError('Could not find identity file in %s.' % ssh_config)
137
138    path = os.path.expanduser(user_config['identityfile'][0])
139    if not os.path.exists(path):
140        raise FileNotFoundError('Specified IdentityFile %s for %s in %s not '
141                                'existing anymore.' % (path, host, ssh_config))
142    return path
143
144
145class SshResults:
146    """Class representing the results from a SSH command to mimic the output
147    of the job.Result class in ACTS.  This is to reduce the changes needed from
148    swapping the ssh connection in ACTS to paramiko.
149
150    Attributes:
151        stdin: The file descriptor to the input channel of the SSH connection.
152        stdout: The file descriptor to the stdout of the SSH connection.
153        stderr: The file descriptor to the stderr of the SSH connection.
154        exit_status: The file descriptor of the SSH command.
155    """
156    def __init__(self, stdin, stdout, stderr, exit_status):
157        self._stdout = stdout.read().decode('utf-8', errors='replace')
158        self._stderr = stderr.read().decode('utf-8', errors='replace')
159        self._exit_status = exit_status.recv_exit_status()
160
161    @property
162    def stdout(self):
163        return self._stdout
164
165    @property
166    def stderr(self):
167        return self._stderr
168
169    @property
170    def exit_status(self):
171        return self._exit_status
172