XTesting/powder-control/powder/ssh.py

   1 #!/usr/bin/env python3
   2 import logging
   3 import os
   4 import pexpect
   5 import re
   6 import time
   7
   8
   9 class SSHConnection:
  10     """A simple ssh/scp wrapper for creating and interacting with ssh sessions via
  11     pexpect.
  12
  13     Args:
  14         ip_address (str): IP address of the node.
  15         username (str): A username with access to the node.
  16         prompt (str) (optional): Expected prompt on the host.
  17
  18     Attributes:
  19         ssh (pexpect child): A handle to a session started by pexpect.spawn()
  20     """
  21
  22     DEFAULT_PROMPT = '\$'
  23
  24     def __init__(self, ip_address, username=None, password=None, prompt=DEFAULT_PROMPT, key=None):
  25         self.prompt = prompt
  26         self.ip_address = ip_address
  27         self.key = key
  28         if username is None:
  29             try:
  30                 self.username = os.environ['USER']
  31             except KeyError:
  32                 logging.error('no USER variable in environment variable and no username provided')
  33                 raise ValueError
  34
  35         if password is None:
  36             try:
  37                 self.password = os.environ['KEYPWORD']
  38             except KeyError:
  39                 logging.info('no ssh key password in environment, assuming unencrypted')
  40                 self.password = password
  41
  42         if key is None:
  43             try:
  44                 self.key = os.environ['CERT']
  45             except KeyError:
  46                 logging.info('no ssh key path provided in environment, assuming not using a ssh key')
  47                 self.key = key
  48
  49     def open(self):
  50         """Opens a connection to `self.ip_address` using `self.username`."""
  51
  52         retry_count = 0
  53         if self.key is None:
  54             cmd = 'ssh -l {} {}'.format(self.username, self.ip_address)
  55         else:
  56             # this is to provide a public key for the ssh session, which might be the most
  57             # commonly use case
  58             cmd = 'ssh -i {} -l {} {}'.format(self.key, self.username, self.ip_address)
  59         while retry_count < 4:
  60             self.ssh = pexpect.spawn(cmd, timeout=5)
  61             self.sshresponse = self.ssh.expect([self.prompt,
  62                                                 'Are you sure you want to continue connecting (yes/no)?',
  63                                                 'Last login',
  64                                                 'Enter passphrase for key.*:',
  65                                                 pexpect.EOF,
  66                                                 pexpect.TIMEOUT])
  67             if self.sshresponse == 0:
  68                 return self
  69             elif self.sshresponse == 1:
  70                 self.ssh.sendline('yes')
  71                 self.sshresponse = self.ssh.expect([self.prompt,
  72                                                     'Enter passphrase for key.*:',
  73                                                     'Permission denied',
  74                                                     pexpect.EOF,
  75                                                     pexpect.TIMEOUT])
  76                 if self.sshresponse == 0:
  77                     return self
  78                 elif self.sshresponse == 1:
  79                     if self.password is None:
  80                         logging.error('failed to login --- ssh key is encrypted but no pword provided')
  81                         raise ValueError
  82
  83                     self.ssh.sendline(self.password)
  84                     self.sshresponse = self.ssh.expect([self.prompt, 'Permission denied', pexpect.EOF, pexpect.TIMEOUT])
  85                     if self.sshresponse == 0:
  86                         return self
  87                     else:
  88                         logging.debug('failed to login --- response: {}'.format(self.sshresponse))
  89                         logging.debug('retry count: {}'.format(retry_count))
  90                 else:
  91                     logging.debug('failed to login --- response: {}'.format(self.sshresponse))
  92                     logging.debug('retry count: {}'.format(retry_count))
  93
  94             elif self.sshresponse == 2:
  95                 # Verify we've connected to the self.ip_address
  96                 self.command('ifconfig | egrep --color=never "inet addr:|inet "', self.prompt)
  97                 self.sshresponse = self.ssh.expect([self.prompt, pexpect.EOF, pexpect.TIMEOUT])
  98                 result = re.search(str(self.ip_address), str(self.ssh.before))
  99                 if result is None:
 100                     logging.debug('not on host with ip {}'.format(self.ip_address))
 101                     logging.debug('retry count: {}'.format(retry_count))
 102                 else:
 103                     return self
 104
 105             elif self.sshresponse == 3:
 106                 if self.password is None:
 107                     logging.error('failed to login --- ssh key is encrypted but no pword provided')
 108                     raise ValueError
 109
 110                 self.ssh.sendline(self.password)
 111                 self.sshresponse = self.ssh.expect([self.prompt, 'Permission denied', pexpect.EOF, pexpect.TIMEOUT])
 112
 113                 if self.sshresponse == 0:
 114                     return self
 115                 else:
 116                     logging.debug('failed to login --- response: {}'.format(self.sshresponse))
 117                     logging.debug('retry count: {}'.format(retry_count))
 118
 119             elif self.sshresponse == 4:
 120                 logging.debug('Unexpected EOF')
 121                 logging.debug('retry count: {}'.format(retry_count))
 122                 logging.debug('ssh.before: ' + str(self.ssh.before))
 123             elif self.sshresponse == 5:
 124                 logging.debug('Unexpected Timeout')
 125                 logging.debug('retry count: {}'.format(retry_count))
 126                 logging.debug('ssh.before: ' + str(self.ssh.before))
 127
 128             time.sleep(1)
 129             retry_count += 1
 130
 131         logging.error('failed to login --- could not connect to host.')
 132         raise ValueError
 133
 134     def command(self, commandline, expectedline=DEFAULT_PROMPT, timeout=5):
 135         """Sends `commandline` to `self.ip_address` and waits for `expectedline`."""
 136         logging.debug(commandline)
 137         self.ssh.sendline(commandline)
 138         self.sshresponse = self.ssh.expect([expectedline, pexpect.EOF, pexpect.TIMEOUT], timeout=timeout)
 139         if self.sshresponse == 0:
 140             pass
 141         elif self.sshresponse == 1:
 142             logging.debug('Unexpected EOF --- Expected: ' + expectedline)
 143             logging.debug('ssh.before: ' + str(self.ssh.before))
 144         elif self.sshresponse == 2:
 145             logging.debug('Unexpected Timeout --- Expected: ' + expectedline)
 146             logging.debug('ssh.before: ' + str(self.ssh.before))
 147
 148         return self.sshresponse
 149
 150     def close(self, timeout):
 151         self.ssh.sendline('exit')
 152         self.sshresponse = self.ssh.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=timeout)
 153         if self.sshresponse == 0:
 154             pass
 155         elif self.sshresponse == 1:
 156             logging.debug('Unexpected Timeout --- Expected: EOF')
 157             logging.debug('ssh.before: ' + str(self.ssh.before))
 158
 159         return self.sshresponse
 160
 161     def copy_from(self, remote_path, local_path):
 162         cmd = 'scp {}@{}:{} {}'.format(self.username, self.ip_address, remote_path, local_path)
 163         return self.copy(cmd)
 164
 165     def copy_to(self, local_path, remote_path):
 166         cmd = 'scp {} {}@{}:{}'.format(local_path, self.username, self.ip_address, remote_path)
 167         return self.copy(cmd)
 168
 169     def copy(self, cmd):
 170         retry_count = 0
 171         logging.debug(cmd)
 172         while retry_count < 10:
 173             scp_spawn = pexpect.spawn(cmd, timeout = 100)
 174             scp_response = scp_spawn.expect(['Are you sure you want to continue connecting (yes/no)?',
 175                                              'Enter passphrase for key.*:',
 176                                              pexpect.EOF,
 177                                              pexpect.TIMEOUT])
 178             if scp_response == 0:
 179                 scp_spawn.sendline('yes')
 180                 scp_response = scp_spawn.expect([self.prompt,
 181                                                  'Enter passphrase for key.*:',
 182                                                  'Permission denied',
 183                                                  pexpect.EOF,
 184                                                  pexpect.TIMEOUT])
 185                 if scp_response == 0:
 186                     return scp_response
 187                 elif scp_response == 1:
 188                     if self.password is None:
 189                         logging.error('failed to scp --- ssh key is encrypted but no pword provided')
 190                         raise ValueError
 191
 192                     scp_spawn.sendline(self.password)
 193                     scp_response = scp_spawn.expect([self.prompt, 'Permission denied', pexpect.EOF, pexpect.TIMEOUT])
 194                     if scp_response == 0:
 195                         return scp_response
 196                     else:
 197                         logging.debug('failed to scp --- response: {}'.format(scp_response))
 198                         logging.debug('retry count: {}'.format(retry_count))
 199                     return scp_response
 200                 else:
 201                     logging.debug('scp failed with scp response {}'.format(scp_response))
 202                     logging.debug('retry count: {}'.format(retry_count))
 203             elif scp_response == 1:
 204                 if self.password is None:
 205                     logging.error('failed to scp --- ssh key is encrypted but no pword provided')
 206                     raise ValueError
 207
 208                 scp_spawn.sendline(self.password)
 209                 scp_response = scp_spawn.expect([self.prompt, 'Permission denied', pexpect.EOF, pexpect.TIMEOUT])
 210                 if scp_response == 0:
 211                     return scp_response
 212                 else:
 213                     logging.debug('failed to scp --- response: {}'.format(scp_response))
 214                     logging.debug('retry count: {}'.format(retry_count))
 215                 return scp_response
 216             elif scp_response == 2:
 217                 logging.debug('copy succeeded')
 218                 return scp_response
 219
 220             time.sleep(1)
 221             retry_count += 1
 222
 223         return scp_response