X-Git-Url: http://dolda2000.com/gitweb/?a=blobdiff_plain;f=pdm%2Fsshsock.py;h=7ee26e899b647ad77356f489c17de446b126db1b;hb=578de289fa6936c1a6b21cc962b5a64d5a3a0e52;hp=feb2f9fcf9c835d9496bb2b56542b62f06fd3eac;hpb=5d0664aa6dde860573fe4bcf8d1b312b7c6c4271;p=pdm.git

diff --git a/pdm/sshsock.py b/pdm/sshsock.py
index feb2f9f..7ee26e8 100644
--- a/pdm/sshsock.py
+++ b/pdm/sshsock.py
@@ -9,9 +9,33 @@ class sshsocket(object):
         if port is not None:
             args += ["-p", str(int(port))]
         args += [host]
-        args += ["python", "-m", "pdm.sshsock", path]
+        args += ["python3", "-m", "pdm.sshsock", path]
         self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, close_fds=True)
         fcntl.fcntl(self.proc.stdout, fcntl.F_SETFL, fcntl.fcntl(self.proc.stdout, fcntl.F_GETFL) | os.O_NONBLOCK)
+        head = self.recv(5)
+        if head != b"SSOCK":
+            raise socket.error("unexpected reply from %s: %r" % (host, head))
+        head = self.recv(1)
+        if head == b"+":
+            buf = b""
+            while True:
+                r = self.recv(1)
+                if r == b"":
+                    raise socket.error("unexpected EOF in SSH socket stream")
+                elif r == b"\n":
+                    break
+                buf += r
+            return
+        elif head == b"-":
+            buf = b""
+            while True:
+                r = self.recv(1)
+                if r in {b"\n", b""}:
+                    break
+                buf += r
+            raise socket.error(buf.decode("utf-8"))
+        else:
+            raise socket.error("unexpected reply from %s: %r" % (host, head))
 
     def close(self):
         if self.proc is not None:
@@ -36,36 +60,43 @@ class sshsocket(object):
         self.close()
 
 def cli():
-    fcntl.fcntl(sys.stdin, fcntl.F_SETFL, fcntl.fcntl(sys.stdin, fcntl.F_GETFL) | os.O_NONBLOCK)
+    fcntl.fcntl(sys.stdin.buffer, fcntl.F_SETFL, fcntl.fcntl(sys.stdin.buffer, fcntl.F_GETFL) | os.O_NONBLOCK)
     sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
     try:
-        sk.connect(sys.argv[1])
-        buf1 = ""
-        buf2 = ""
+        try:
+            sk.connect(sys.argv[1])
+        except socket.error as err:
+            sys.stdout.write("SSOCK-connect: %s\n" % err)
+            sys.stdout.flush()
+            return
+        sys.stdout.write("SSOCK+\n")
+        sys.stdout.flush()
+        buf1 = b""
+        buf2 = b""
         while True:
             wfd = []
             if buf1: wfd.append(sk)
-            if buf2: wfd.append(sys.stdout)
-            rfd, wfd, efd = select.select([sk, sys.stdin], wfd, [])
+            if buf2: wfd.append(sys.stdout.buffer)
+            rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, [])
             if sk in rfd:
                 ret = sk.recv(65536)
-                if ret == "":
+                if ret == b"":
                     break
                 else:
                     buf2 += ret
-            if sys.stdin in rfd:
-                ret = sys.stdin.read()
-                if ret == "":
+            if sys.stdin.buffer in rfd:
+                ret = sys.stdin.buffer.read()
+                if ret == b"":
                     break
                 else:
                     buf1 = ret
             if sk in wfd:
                 ret = sk.send(buf1)
                 buf1 = buf1[ret:]
-            if sys.stdout in wfd:
-                sys.stdout.write(buf2)
-                sys.stdout.flush()
-                buf2 = ""
+            if sys.stdout.buffer in wfd:
+                sys.stdout.buffer.write(buf2)
+                sys.stdout.buffer.flush()
+                buf2 = b""
     finally:
         sk.close()