Merge pull request #712 from kbriggs/zsh-login-prompt

Set prompt correctly for zsh

3 files changed, 48 insertions, 11 deletions

diff --git a/pexpect/pxssh.py b/pexpect/pxssh.py
index bfefc7a..742f59e 100644
--- a/pexpect/pxssh.py
+++ b/pexpect/pxssh.py
@@ -143,6 +143,7 @@ class pxssh (spawn):
         # used to set shell command-line prompt to UNIQUE_PROMPT.
         self.PROMPT_SET_SH = r"PS1='[PEXPECT]\$ '"
         self.PROMPT_SET_CSH = r"set prompt='[PEXPECT]\$ '"
+        self.PROMPT_SET_ZSH = "prompt restore;\nPS1='[PEXPECT]%(!.#.$) '"
         self.SSH_OPTS = (" -o 'PubkeyAuthentication=no'")
 # Disabling host key checking, makes you vulnerable to MITM attacks.
 #                + " -o 'StrictHostKeyChecking=no'"
@@ -529,8 +530,11 @@ class pxssh (spawn):
         if i == 0: # csh-style
             self.sendline(self.PROMPT_SET_CSH)
             i = self.expect([TIMEOUT, self.PROMPT], timeout=10)
-            if i == 0:
-                return False
+            if i == 0: # zsh-style
+                self.sendline(self.PROMPT_SET_ZSH)
+                i = self.expect([TIMEOUT, self.PROMPT], timeout=10)
+                if i == 0:
+                    return False
         return True
 
 # vi:ts=4:sw=4:expandtab:ft=python:
diff --git a/tests/fakessh/ssh b/tests/fakessh/ssh
index 74ffe20..e794416 100755
--- a/tests/fakessh/ssh
+++ b/tests/fakessh/ssh
@@ -9,19 +9,14 @@ if not PY3:
     input = raw_input
 
 ssh_usage = "usage: ssh [-2qV] [-c cipher_spec] [-l login_name]\r\n" \
-          + "           hostname"
+          + "           hostname [shell]"
 
 cipher_valid_list = ['aes128-ctr', 'aes192-ctr', 'aes256-ctr', 'arcfour256', 'arcfour128', \
     'aes128-cbc','3des-cbc','blowfish-cbc','cast128-cbc','aes192-cbc', \
     'aes256-cbc','arcfour']
 
 try:
-    server = sys.argv[-1]
-    if server == 'noserver':
-        print('No route to host')
-        sys.exit(1)
-
-    elif len(sys.argv) < 2:
+    if len(sys.argv) < 2:
         print(ssh_usage)
         sys.exit(1)
 
@@ -29,7 +24,7 @@ try:
     cipher_list = []
     fullCmdArguments = sys.argv
     argumentList = fullCmdArguments[1:]
-    unixOptions = "2qVc:l"
+    unixOptions = "2qVc:l:"
     arguments, values = getopt.getopt(argumentList, unixOptions)
     for currentArgument, currentValue in arguments:
         if currentArgument in ("-2"):
@@ -45,6 +40,14 @@ try:
                     print("Unknown cipher type '" + str(cipher_item) + "'")
                     sys.exit(1)
 
+    server = values[0]
+    if server == 'noserver':
+        print('No route to host')
+        sys.exit(1)
+
+    shell = 'bash'
+    if len(values) > 1:
+        shell = values[1]
 
 except Exception as e:
     print(ssh_usage)
@@ -62,7 +65,13 @@ prompt = "$"
 while True:
     cmd = input(prompt)
     if cmd.startswith('PS1='):
-        prompt = eval(cmd[4:]).replace(r'\$', '$')
+        if shell == 'bash':
+            prompt = eval(cmd[4:]).replace(r'\$', '$')
+        elif shell == 'zsh':
+            prompt = eval(cmd[4:]).replace('%(!.#.$)', '$')
+    elif cmd.startswith('set prompt='):
+        if shell.endswith('csh'):
+            prompt = eval(cmd[11:]).replace(r'\$', '$')
     elif cmd == 'ping':
         print('pong')
     elif cmd.startswith('ls'):
diff --git a/tests/test_pxssh.py b/tests/test_pxssh.py
index ba700c8..f768d22 100644
--- a/tests/test_pxssh.py
+++ b/tests/test_pxssh.py
@@ -276,5 +276,29 @@ class PxsshTestCase(SSHTestBase):
         else:
             assert False, 'should have raised exception, pxssh.ExceptionPxssh'
 
+    def test_login_bash(self):
+        ssh = pxssh.pxssh()
+        result = ssh.login('server bash', 'me', password='s3cret')
+        ssh.sendline('ping')
+        ssh.expect('pong', timeout=10)
+        assert ssh.prompt(timeout=10)
+        ssh.logout()
+
+    def test_login_zsh(self):
+        ssh = pxssh.pxssh()
+        result = ssh.login('server zsh', 'me', password='s3cret')
+        ssh.sendline('ping')
+        ssh.expect('pong', timeout=10)
+        assert ssh.prompt(timeout=10)
+        ssh.logout()
+
+    def test_login_tcsh(self):
+        ssh = pxssh.pxssh()
+        result = ssh.login('server tcsh', 'me', password='s3cret')
+        ssh.sendline('ping')
+        ssh.expect('pong', timeout=10)
+        assert ssh.prompt(timeout=10)
+        ssh.logout()
+
 if __name__ == '__main__':
     unittest.main()