Refining the code.

1 files changed, 64 insertions, 31 deletions

diff --git a/pexpect/pxssh.py b/pexpect/pxssh.py
index d228f03..4a87df7 100644
--- a/pexpect/pxssh.py
+++ b/pexpect/pxssh.py
@@ -1,5 +1,5 @@
 """This class extends pexpect.spawn to specialize setting up SSH connections.
-This adds methods for login, logout, and expecting the prompt.
+This adds methods for login, logout, and expecting the shell prompt.
 $Id$
 """
 from pexpect import *
@@ -13,24 +13,40 @@ class ExceptionPxssh(ExceptionPexpect):
 
 class pxssh (spawn):
     """This class extends pexpect.spawn to specialize setting up SSH connections.
-    This adds methods for login, logout, and expecting the prompt.
-    It does various hacky things to handle any situation in the SSH login process.
-    For example, if the session is your first login, then it automatically
-    accepts the certificate; or if you have public key authentication setup
-    and you don't need a password then this is OK too.
+    This adds methods for login, logout, and expecting the shell prompt.
+    It does various tricky things to handle many situations in the SSH login process.
+    For example, if the session is your first login, then pxssh automatically
+    accepts the remote certificate; or if you have public key authentication setup
+    then pxssh won't wait for the password prompt.
 
-    Example usage that runs 'ls -l' on server and prints the result:
+    pxssh uses the shell prompt to synchronize output from the remote host. In
+    order to make this more robust it sets the shell prompt to something more
+    unique than just $ or #. This should work on most Borne/Bash or Csh style
+    shells.
+
+    Example usage that runs a few commands on a remote server and prints the result:
         import pxssh
-        s = pxssh.pxssh()
-        if not s.login ('localhost', 'myusername', 'mypassword'):
-            print "SSH session failed on login."
-            print str(s)
-        else:
-            print "SSH session login successful"
+        import getpass
+        try:                                                            
+            s = pxssh.pxssh()
+            hostname = raw_input('hostname: ')
+            username = raw_input('username: ')
+            password = getpass.getpass('password: ')
+            s.login (hostname, username, password)
+            s.sendline ('uptime')  # run a command
+            s.prompt()             # match the prompt
+            print s.before         # print everything before the prompt.
             s.sendline ('ls -l')
-            s.prompt()           # match the prompt
-            print s.before     # print everything before the prompt.
+            s.prompt()
+            print s.before
+            s.sendline ('df')
+            s.prompt()
+            print s.before
             s.logout()
+        except pxssh.ExceptionPxssh, e:
+            print "pxssh failed on login."
+            print str(e)
+
     """
 
     def __init__ (self, timeout=30, maxread=2000, searchwindowsize=None, logfile=None, env=None):
@@ -43,15 +59,15 @@ class pxssh (spawn):
         # command. So if we make the set command slightly different than the regex
         # we eliminate the problem. To make the set command different we add a
         # backslash in front of $. The $ doesn't need to be escaped, but it doesn't
-        # hurt and serves to make the set command different than the regex.
+        # hurt and serves to make the set prompt command different than the regex.
         self.PROMPT = "\[PEXPECT\][\$\#] " # used to match the command-line prompt.
         # used to set shell command-line prompt to something more unique.
         self.PROMPT_SET_SH = "PS1='[PEXPECT]\$ '"
         self.PROMPT_SET_CSH = "set prompt='[PEXPECT]\$ '"
 
     ### TODO: This is getting messy and I'm pretty sure this isn't perfect.
-    ### TODO: I need to draw a better flow chart for this.
-    def login (self,server,username,password='',terminal_type='ansi',original_prompts=r"][#$]|~[#$]|bash.*?[#$]|[#$] ",login_timeout=10):
+    ### TODO: I need to draw a flow chart for this.
+    def login (self,server,username,password='',terminal_type='ansi',original_prompts=r"][#$]|~[#$]|bash.*?[#$]|[#$] ",login_timeout=10,port=None):
         """This logs the user into the given server. By default the prompt is
         rather optimistic and should be considered more of an example. It's
         better to try to match the prompt as exactly as possible to prevent
@@ -62,7 +78,10 @@ class pxssh (spawn):
         we could not match it. We still try to reset the prompt to something
         more unique. If that still fails then we return False.
         """
-        cmd = "ssh -l %s %s" % (username, server)
+        if port is None:
+            cmd = "ssh -l %s %s" % (username, server)
+        else:
+            cmd = "ssh -p %s -l %s %s" % (str(port),username,server)
         spawn.__init__(self, cmd, timeout=login_timeout)
         #, "(?i)no route to host"])
         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"])
@@ -82,7 +101,7 @@ class pxssh (spawn):
             #return False
             raise ExceptionPxssh ('Weird error. Got "are you sure" prompt twice.')
         elif i==1: # can occur if you have a public key pair set to authenticate. 
-            ### TODO: May NOT be OK if expect() matched a false prompt.
+            ### TODO: May NOT be OK if expect() got tricked and matched a false prompt.
             pass
         elif i==2: # password prompt again
             # For incorrect passwords, some ssh servers will
@@ -101,8 +120,12 @@ class pxssh (spawn):
             #return False
             raise ExceptionPxssh ('Weird error. Got "terminal type" prompt twice.')
         elif i==5: # Timeout
-            # This is tricky... presume that we are at the command-line prompt.
-            # It may be that the prompt was so weird that we couldn't match it.
+            # This is tricky... I presume that we are at the command-line prompt.
+            # It may be that the shell prompt was so weird that we couldn't match it.
+            # Or it may be that we couldn't log in for some other reason.
+            # I can't be sure, but it's safe to guess that we did loging because
+            # if I presume wrong and we are not logged in then this should
+            # be caught later when I try to set the shell prompt.
             pass
         elif i==6: # Connection closed by remote host
             self.close()
@@ -112,15 +135,16 @@ class pxssh (spawn):
             self.close()
             #return False
             raise ExceptionPxssh ('unexpected login response')
-        # We appear to be in -- reset prompt to something more unique.
+        # We appear to be in -- set shell prompt to something more unique.
         if not self.set_unique_prompt():
             self.close()
             #return False
-            raise ExceptionPxssh ('could not set prompt')
+            raise ExceptionPxssh ('could not set shell prompt\n'+self.before)
         return True
 
     def logout (self):
-        """This sends exit. If there are stopped jobs then this sends exit twice.
+        """This sends exit to the remote shell.
+        If there are stopped jobs then this sends exit twice.
         """
         self.sendline("exit")
         index = self.expect([EOF, "(?i)there are stopped jobs"])
@@ -129,7 +153,9 @@ class pxssh (spawn):
             self.expect(EOF)
 
     def prompt (self, timeout=20):
-        """This expects the prompt. This returns True if the prompt was matched.
+        """This matches the shell prompt.
+        This is little more than a short-cut to the expect() method.
+        This returns True if the shell prompt was matched.
         This returns False if there was a timeout.
         """
         i = self.expect([self.PROMPT, TIMEOUT], timeout=timeout)
@@ -137,12 +163,19 @@ class pxssh (spawn):
             return False
         return True
         
-    def set_unique_prompt (self, optional_prompt=None):
-        """This attempts to reset the shell prompt to something more unique.
-            This makes it easier to match unambiguously.
+    def set_unique_prompt (self):
+        """This sets the remote shell prompt to something more unique than
+        # or $. This makes it easier for the prompt() method to match the shell
+        prompt unambiguously. This method is called automatically by the
+        login() method, but you may want to call it manually if you somehow
+        reset the shell prompt. For example, if you 'su' to a different user
+        then you will need to manually reset the prompt. This sends shell
+        commands to the remote host to set the prompt, so this assumes the
+        remote host is ready to receive commands.
+
+        You can also set your own prompt and set the PROMPT attribute
+        to a regular expression that matches it.
         """
-        if optional_prompt is not None:
-            self.PROMPT = optional_prompt
         self.sendline (self.PROMPT_SET_SH) # sh-style
         i = self.expect ([TIMEOUT, self.PROMPT], timeout=10)
         if i == 0: # csh-style