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

Set prompt correctly for zsh

pexpect/pxssh.py

@@ -143,6 +143,7 @@ def __init__ (self, timeout=30, maxread=2000, searchwindowsize=None,
         # 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 @@ def set_unique_prompt(self):
         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:

tests/fakessh/ssh

@@ -9,27 +9,22 @@ 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)
 
     cipher = ''
     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'):

tests/test_pxssh.py

@@ -276,5 +276,29 @@ def test_failed_custom_ssh_cmd(self):
         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()