Merge branch 'master' into python3
[pdm.git] / pdm / sshsock.py
... / ...
CommitLineData
1import sys, os
2import subprocess, socket, fcntl, select
3
4class 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
62def 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 return
71 sys.stdout.write("SSOCK+\n")
72 buf1 = b""
73 buf2 = b""
74 while True:
75 wfd = []
76 if buf1: wfd.append(sk)
77 if buf2: wfd.append(sys.stdout.buffer)
78 rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, [])
79 if sk in rfd:
80 ret = sk.recv(65536)
81 if ret == b"":
82 break
83 else:
84 buf2 += ret
85 if sys.stdin.buffer in rfd:
86 ret = sys.stdin.buffer.read()
87 if ret == b"":
88 break
89 else:
90 buf1 = ret
91 if sk in wfd:
92 ret = sk.send(buf1)
93 buf1 = buf1[ret:]
94 if sys.stdout.buffer in wfd:
95 sys.stdout.buffer.write(buf2)
96 sys.stdout.buffer.flush()
97 buf2 = b""
98 finally:
99 sk.close()
100
101if __name__ == "__main__":
102 cli()