Made relevant Python3 changes.
[pdm.git] / pdm / sshsock.py
1 import sys, os
2 import subprocess, socket, fcntl, select
3
4 class sshsocket(object):
5     def __init__(self, host, path, user = None, port = None):
6         args = ["ssh"]
7         if user is not None:
8             args += ["-u", str(user)]
9         if port is not None:
10             args += ["-p", str(int(port))]
11         args += [host]
12         args += ["python3", "-m", "pdm.sshsock", path]
13         self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, close_fds=True)
14         fcntl.fcntl(self.proc.stdout, fcntl.F_SETFL, fcntl.fcntl(self.proc.stdout, fcntl.F_GETFL) | os.O_NONBLOCK)
15         head = self.recv(5)
16         if head != b"SSOCK":
17             raise socket.error("unexpected reply from %s: %r" % (host, head))
18         head = self.recv(1)
19         if head == b"+":
20             buf = b""
21             while True:
22                 r = self.recv(1)
23                 if r == b"":
24                     raise socket.error("unexpected EOF in SSH socket stream")
25                 elif r == b"\n":
26                     break
27                 buf += r
28             return
29         elif head == b"-":
30             buf = b""
31             while True:
32                 r = self.recv(1)
33                 if r in {b"\n", b""}:
34                     break
35                 buf += r
36             raise socket.error(buf.decode("utf-8"))
37         else:
38             raise socket.error("unexpected reply from %s: %r" % (host, head))
39
40     def close(self):
41         if self.proc is not None:
42             self.proc.stdin.close()
43             self.proc.stdout.close()
44             self.proc.wait()
45             self.proc = None
46
47     def send(self, data, flags = 0):
48         self.proc.stdin.write(data)
49         return len(data)
50
51     def recv(self, buflen, flags = 0):
52         if (flags & socket.MSG_DONTWAIT) == 0:
53             select.select([self.proc.stdout], [], [])
54         return self.proc.stdout.read(buflen)
55
56     def fileno(self):
57         return self.proc.stdout.fileno()
58
59     def __del__(self):
60         self.close()
61
62 def cli():
63     fcntl.fcntl(sys.stdin.buffer, fcntl.F_SETFL, fcntl.fcntl(sys.stdin.buffer, fcntl.F_GETFL) | os.O_NONBLOCK)
64     sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
65     try:
66         try:
67             sk.connect(sys.argv[1])
68         except socket.error as err:
69             sys.stdout.write("SSOCK-connect: %s\n" % err)
70             sys.stdout.flush()
71             return
72         sys.stdout.write("SSOCK+\n")
73         sys.stdout.flush()
74         buf1 = b""
75         buf2 = b""
76         while True:
77             wfd = []
78             if buf1: wfd.append(sk)
79             if buf2: wfd.append(sys.stdout.buffer)
80             rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, [])
81             if sk in rfd:
82                 ret = sk.recv(65536)
83                 if ret == b"":
84                     break
85                 else:
86                     buf2 += ret
87             if sys.stdin.buffer in rfd:
88                 ret = sys.stdin.buffer.read()
89                 if ret == b"":
90                     break
91                 else:
92                     buf1 = ret
93             if sk in wfd:
94                 ret = sk.send(buf1)
95                 buf1 = buf1[ret:]
96             if sys.stdout.buffer in wfd:
97                 sys.stdout.buffer.write(buf2)
98                 sys.stdout.buffer.flush()
99                 buf2 = b""
100     finally:
101         sk.close()
102
103 if __name__ == "__main__":
104     cli()