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
24logging.getLogger("paramiko").setLevel(logging.WARNING)
25
26
27def get_private_key(ip_address, ssh_config):
28    """Tries to load various ssh key types.
29
30    Args:
31        ip_address: IP address of ssh server.
32        ssh_config: ssh_config location for the ssh server.
33    Returns:
34        The ssh private key
35    """
36    exceptions = []
37    try:
38        logging.debug('Trying to load SSH key type: ed25519')
39        return paramiko.ed25519key.Ed25519Key(
40            filename=get_ssh_key_for_host(ip_address, ssh_config))
41    except paramiko.SSHException as e:
42        exceptions.append(e)
43        logging.debug('Failed loading SSH key type: ed25519')
44
45    try:
46        logging.debug('Trying to load SSH key type: rsa')
47        return paramiko.RSAKey.from_private_key_file(
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: rsa')
52
53    raise Exception('No valid ssh key type found', exceptions)
54
55
56@backoff.on_exception(
57    backoff.constant,
58    (paramiko.ssh_exception.SSHException,
59     paramiko.ssh_exception.AuthenticationException, socket.timeout,
60     socket.error, ConnectionRefusedError, ConnectionResetError),
61    interval=1.5,
62    max_tries=4)
63def create_ssh_connection(ip_address,
64                          ssh_username,
65                          ssh_config,
66                          connect_timeout=30):
67    """Creates and ssh connection to a Fuchsia device
68
69    Args:
70        ip_address: IP address of ssh server.
71        ssh_username: Username for ssh server.
72        ssh_config: ssh_config location for the ssh server.
73        connect_timeout: Timeout value for connecting to ssh_server.
74
75    Returns:
76        A paramiko ssh object
77    """
78    ssh_key = get_private_key(ip_address=ip_address, ssh_config=ssh_config)
79    ssh_client = paramiko.SSHClient()
80    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
81    ssh_client.connect(hostname=ip_address,
82                       username=ssh_username,
83                       allow_agent=False,
84                       pkey=ssh_key,
85                       timeout=connect_timeout,
86                       banner_timeout=200)
87    return ssh_client
88
89
90def ssh_is_connected(ssh_client):
91    """Checks to see if the SSH connection is alive.
92    Args:
93        ssh_client: A paramiko SSH client instance.
94    Returns:
95          True if connected, False or None if not connected.
96    """
97    return ssh_client and ssh_client.get_transport().is_active()
98
99
100def get_ssh_key_for_host(host, ssh_config_file):
101    """Gets the SSH private key path from a supplied ssh_config_file and the
102       host.
103    Args:
104        host (str): The ip address or host name that SSH will connect to.
105        ssh_config_file (str): Path to the ssh_config_file that will be used
106            to connect to the host.
107
108    Returns:
109        path: A path to the private key for the SSH connection.
110    """
111    ssh_config = paramiko.SSHConfig()
112    user_config_file = os.path.expanduser(ssh_config_file)
113    if os.path.exists(user_config_file):
114        with open(user_config_file) as f:
115            ssh_config.parse(f)
116    user_config = ssh_config.lookup(host)
117
118    if 'identityfile' not in user_config:
119        raise ValueError('Could not find identity file in %s.' % ssh_config)
120
121    path = os.path.expanduser(user_config['identityfile'][0])
122    if not os.path.exists(path):
123        raise FileNotFoundError('Specified IdentityFile %s for %s in %s not '
124                                'existing anymore.' % (path, host, ssh_config))
125    return path
126
127
128class SshResults:
129    """Class representing the results from a SSH command to mimic the output
130    of the job.Result class in ACTS.  This is to reduce the changes needed from
131    swapping the ssh connection in ACTS to paramiko.
132
133    Attributes:
134        stdin: The file descriptor to the input channel of the SSH connection.
135        stdout: The file descriptor to the stdout of the SSH connection.
136        stderr: The file descriptor to the stderr of the SSH connection.
137        exit_status: The file descriptor of the SSH command.
138    """
139    def __init__(self, stdin, stdout, stderr, exit_status):
140        self._stdout = stdout.read().decode('utf-8', errors='replace')
141        self._stderr = stderr.read().decode('utf-8', errors='replace')
142        self._exit_status = exit_status.recv_exit_status()
143
144    @property
145    def stdout(self):
146        return self._stdout
147
148    @property
149    def stderr(self):
150        return self._stderr
151
152    @property
153    def exit_status(self):
154        return self._exit_status
155