pcucontrol/transports/ssh/pxssh.py

   1 from pexpect import *
   2 #import os, sys, getopt, shutil
   3
   4 class pxssh (spawn):
   5     """This class extends pexpect.spawn to specialize setting up SSH connections.
   6     This adds methods for login, logout, and expecting the prompt.
   7     It does various hacky things to handle any situation in the SSH login process.
   8     For example, if the session is your first login, then it automatically
   9     accepts the certificate; or if you have public key authentication setup
  10     and you don't need a password then this is OK too.
  11
  12     Example usage that runs 'ls -l' on server and prints the result:
  13         import pxssh
  14         s = pxssh.pxssh()
  15         if not s.login ('localhost', 'myusername', 'mypassword'):
  16             print "SSH session failed on login."
  17             print str(s)
  18         else:
  19             print "SSH session login successful"
  20             s.sendline ('ls -l')
  21             s.prompt()           # match the prompt
  22             print s.before     # print everything before the prompt.
  23             s.logout()
  24     """
  25
  26     def __init__ (self):
  27         # SUBTLE HACK ALERT!
  28         # Note that the command to set the prompt uses a slightly different string
  29         # than the regular expression to match it. This is because when you set the
  30         # prompt the command will echo back, but we don't want to match the echoed
  31         # command. So if we make the set command slightly different than the regex
  32         # we eliminate the problem. To make the set command different we add a
  33         # backslash in front of $. The $ doesn't need to be escaped, but it doesn't
  34         # hurt and serves to make the set command different than the regex.
  35         self.PROMPT = "\[PEXPECT\][\$\#] " # used to match the command-line prompt.
  36         # used to set shell command-line prompt to something more unique.
  37         self.PROMPT_SET_SH = "PS1='[PEXPECT]\$ '"
  38         self.PROMPT_SET_CSH = "set prompt='[PEXPECT]\$ '"
  39
  40     ### TODO: This is getting messy and I'm pretty sure this isn't perfect.
  41     ### TODO: I need to draw a better flow chart for this.
  42     def login (self,server,username,password='',ssh_options="",terminal_type='ansi',original_prompts=r"][#$]|~[#$]|-bash.*?[#$]|[#$] ",login_timeout=10):
  43         """This logs the user into the given server. By default the prompt is
  44         rather optimistic and should be considered more of an example. It's
  45         better to try to match the prompt as exactly as possible to prevent
  46         any false matches by server strings such as a "Message Of The Day" or
  47         something. The closer you can make the original_prompt match your real prompt
  48         then the better. A timeout will not necessarily cause the login to fail.
  49         In the case of a timeout we assume that the prompt was so weird that
  50         we could not match it. We still try to reset the prompt to something
  51         more unique. If that still fails then we return False.
  52         """
  53         cmd = "ssh %s -l %s %s" % (ssh_options, username, server)
  54         #print cmd
  55         spawn.__init__(self, cmd, timeout=login_timeout)
  56         #, "(?i)no route to host"])
  57         i = self.expect(["(?i)are you sure you want to continue connecting", original_prompts, "(?i)password", "(?i)permission denied", "(?i)terminal type", TIMEOUT, "(?i)connection closed by remote host"])
  58         if i==0: # New certificate -- always accept it. This is what you if SSH does not have the remote host's public key stored in the cache.
  59             self.sendline("yes")
  60             i = self.expect(["(?i)are you sure you want to continue connecting", original_prompts, "(?i)password", "(?i)permission denied", "(?i)terminal type", TIMEOUT])
  61         if i==2: # password
  62             self.sendline(password)
  63             i = self.expect(["(?i)are you sure you want to continue connecting", original_prompts, "(?i)password", "(?i)permission denied", "(?i)terminal type", TIMEOUT])
  64         if i==4:
  65             self.sendline(terminal_type)
  66             i = self.expect(["(?i)are you sure you want to continue connecting", original_prompts, "(?i)password", "(?i)permission denied", "(?i)terminal type", TIMEOUT])
  67
  68         if i==0:
  69             # This is weird. This should not happen twice in a row.
  70             self.close()
  71             return False
  72         elif i==1: # can occur if you have a public key pair set to authenticate.
  73             ### TODO: May NOT be OK if expect() matched a false prompt.
  74             pass
  75         elif i==2: # password prompt again
  76             # For incorrect passwords, some ssh servers will
  77             # ask for the password again, others return 'denied' right away.
  78             # If we get the password prompt again then this means
  79             # we didn't get the password right the first time.
  80             self.close()
  81             return False
  82         elif i==3: # permission denied -- password was bad.
  83             self.close()
  84             return False
  85         elif i==4: # terminal type again? WTF?
  86             self.close()
  87             return False
  88         elif i==5: # Timeout
  89             # This is tricky... presume that we are at the command-line prompt.
  90             # It may be that the prompt was so weird that we couldn't match it.
  91             pass
  92         elif i==6: # Connection closed by remote host
  93             self.close()
  94             return False
  95         else: # Unexpected
  96             self.close()
  97             return False
  98         # We appear to be in -- reset prompt to something more unique.
  99         #if not self.set_unique_prompt():
 100         #    print "couldn't reset prompt"
 101         #    self.close()
 102         #    return False
 103         return True
 104
 105     def logout (self):
 106         """This sends exit. If there are stopped jobs then this sends exit twice.
 107         """
 108         self.sendline("exit")
 109         index = self.expect([EOF, "(?i)there are stopped jobs"])
 110         if index==1:
 111             self.sendline("exit")
 112             self.expect(EOF)
 113
 114     def prompt (self, timeout=20):
 115         """This expects the prompt. This returns True if the prompt was matched.
 116         This returns False if there was a timeout.
 117         """
 118         i = self.expect([self.PROMPT, TIMEOUT], timeout=timeout)
 119         if i==1:
 120             return False
 121         return True
 122
 123     def set_unique_prompt (self, optional_prompt=None):
 124         """This attempts to reset the shell prompt to something more unique.
 125             This makes it easier to match unambiguously.
 126         """
 127         if optional_prompt is not None:
 128             self.prompt = optional_prompt
 129         self.sendline (self.PROMPT_SET_SH) # sh-style
 130         i = self.expect ([TIMEOUT, self.PROMPT], timeout=10)
 131         if i == 0: # csh-style
 132             self.sendline (self.PROMPT_SET_CSH)
 133             i = self.expect ([TIMEOUT, self.PROMPT], timeout=10)
 134             if i == 0:
 135                 return 0
 136         return 1
 137