# ssh.py - Interact with remote SSH servers

#

# Copyright 2019 Gregory Szorc <gregory.szorc@gmail.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

# no-check-code because Python 3 native.

import socket

import time

import warnings

from cryptography.utils import (

    CryptographyDeprecationWarning,

)

import paramiko

def wait_for_ssh(hostname, port, timeout=60, username=None, key_filename=None):

    """Wait for an SSH server to start on the specified host and port."""

    class IgnoreHostKeyPolicy(paramiko.MissingHostKeyPolicy):

        def missing_host_key(self, client, hostname, key):

            return

    end_time = time.time() + timeout

    # paramiko triggers a CryptographyDeprecationWarning in the cryptography

    # package. Let's suppress

    with warnings.catch_warnings():

        warnings.filterwarnings('ignore',

                                category=CryptographyDeprecationWarning)

        while True:

            client = paramiko.SSHClient()

            client.set_missing_host_key_policy(IgnoreHostKeyPolicy())

            try:

                client.connect(hostname, port=port, username=username,

                               key_filename=key_filename,

                               timeout=5.0, allow_agent=False,

                               look_for_keys=False)

                return client

            except socket.error:

                pass

            except paramiko.AuthenticationException:

                raise

            except paramiko.SSHException:

                pass

            if time.time() >= end_time:

                raise Exception('Timeout reached waiting for SSH')

            time.sleep(1.0)

def exec_command(client, command):

    """exec_command wrapper that combines stderr/stdout and returns channel"""

    chan = client.get_transport().open_session()

    chan.exec_command(command)

    chan.set_combine_stderr(True)

    stdin = chan.makefile('wb', -1)

    stdout = chan.makefile('r', -1)

    return chan, stdin, stdout