1 # ssh.py - Interact with remote SSH servers

     2 #

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

     4 #

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

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

     7 

     8 # no-check-code because Python 3 native.

     9 

    10 import socket

    11 import time

    12 import warnings

    13 

    14 from cryptography.utils import (

    15     CryptographyDeprecationWarning,

    16 )

    17 import paramiko

    18 

    19 

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

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

    22     class IgnoreHostKeyPolicy(paramiko.MissingHostKeyPolicy):

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

    24             return

    25 

    26     end_time = time.time() + timeout

    27 

    28     # paramiko triggers a CryptographyDeprecationWarning in the cryptography

    29     # package. Let's suppress

    30     with warnings.catch_warnings():

    31         warnings.filterwarnings('ignore',

    32                                 category=CryptographyDeprecationWarning)

    33 

    34         while True:

    35             client = paramiko.SSHClient()

    36             client.set_missing_host_key_policy(IgnoreHostKeyPolicy())

    37             try:

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

    39                                key_filename=key_filename,

    40                                timeout=5.0, allow_agent=False,

    41                                look_for_keys=False)

    42 

    43                 return client

    44             except socket.error:

    45                 pass

    46             except paramiko.AuthenticationException:

    47                 raise

    48             except paramiko.SSHException:

    49                 pass

    50 

    51             if time.time() >= end_time:

    52                 raise Exception('Timeout reached waiting for SSH')

    53 

    54             time.sleep(1.0)

    55 

    56 

    57 def exec_command(client, command):

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

    59     chan = client.get_transport().open_session()

    60 

    61     chan.exec_command(command)

    62     chan.set_combine_stderr(True)

    63 

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

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

    66 

    67     return chan, stdin, stdout